"""Auto-fit solvers: pick n_ctx, quant, or GPU setup from a VRAM budget.

These call :func:`vramcalc.report.estimate` in a loop — no new VRAM math — so
they stay consistent with the manual calculator. All three are pure and
unit-testable.

  max_context   — largest n_ctx that fits the GPU budget at a given quant.
  best_quant    — highest-bpw quant that fits a given n_ctx (plus fallback).
  min_gpu_setup — smallest subset of the user's GPUs that fits (largest-first).
"""

from __future__ import annotations

from dataclasses import dataclass, field, replace
from typing import Callable

from .quant import QUANT_BPW
from .report import ModelArch, Inputs, estimate


def _total_vram_bytes(inp: Inputs) -> int:
    return int(sum(g * (1 << 30) for g in inp.gpu_vram_gb))


def _fits(bd) -> bool:
    return bd.gpu is not None and bd.gpu.all_fit


@dataclass
class MaxContextResult:
    n_ctx: int = 0
    total_bytes: float = 0.0
    fits: bool = False
    note: str = ""


def max_context(
    arch: ModelArch,
    inp: Inputs,
    *,
    step: int = 512,
    max_ctx: int | None = None,
) -> MaxContextResult:
    """Binary-search the largest n_ctx whose estimate fits the GPU budget.

    Search bounds: [step … max_ctx]. Default max_ctx = max(training_ctx, n_ctx)
    * 4 to allow YaRN-extended contexts without an unbounded loop. Returns the
    largest fitting n_ctx (rounded down to a multiple of `step`) and the
    resulting total; ``fits`` is False if even the minimum doesn't fit.
    """
    budget = _total_vram_bytes(inp)
    hi = max_ctx or max(arch.training_ctx, inp.n_ctx, step) * 4
    hi = max(hi, step)
    lo = step

    # quick check: does the minimum fit at all?
    bd_min = estimate(arch, replace(inp, n_ctx=lo))
    if not _fits(bd_min):
        return MaxContextResult(
            n_ctx=lo, total_bytes=bd_min.total_bytes, fits=False,
            note=f"Even n_ctx={lo} doesn't fit {budget/1e9:.1f} GiB budget.",
        )

    # does the maximum fit? (small models on big GPUs)
    bd_hi = estimate(arch, replace(inp, n_ctx=hi))
    if _fits(bd_hi):
        return MaxContextResult(
            n_ctx=hi, total_bytes=bd_hi.total_bytes, fits=True,
            note=f"Even n_ctx={hi} fits; try a higher cap.",
        )

    # binary search the boundary
    best = lo
    best_total = bd_min.total_bytes
    while lo <= hi:
        mid = (lo + hi) // 2
        # snap to step
        mid = (mid // step) * step
        if mid < step:
            mid = step
        bd = estimate(arch, replace(inp, n_ctx=mid))
        if _fits(bd):
            best = mid
            best_total = bd.total_bytes
            lo = mid + step
        else:
            hi = mid - step
    return MaxContextResult(n_ctx=best, total_bytes=best_total, fits=True)


@dataclass
class BestQuantResult:
    quant: str = ""
    bpw: float = 0.0
    total_bytes: float = 0.0
    fits: bool = False
    fallback_quant: str = ""
    fallback_bpw: float = 0.0
    fallback_total: float = 0.0
    note: str = ""


def best_quant(arch: ModelArch, inp: Inputs) -> BestQuantResult:
    """Highest-bpw quant that fits the GPU budget at the current n_ctx.

    Iterates quants from highest to lowest bpw; the first that fits is the
    recommendation, the next-lower that also fits is the safe fallback.
    """
    ordered = sorted(QUANT_BPW.items(), key=lambda kv: kv[1], reverse=True)
    res = BestQuantResult()
    found_fit = None
    fallback = None
    for q, bpw in ordered:
        bd = estimate(arch, replace(inp, quant=q))
        fits = _fits(bd)
        if fits and found_fit is None:
            found_fit = (q, bpw, bd.total_bytes)
        elif fits and found_fit is not None and fallback is None:
            fallback = (q, bpw, bd.total_bytes)
            break  # first lower-bpw fit is the fallback
        elif not fits and found_fit is not None and fallback is None:
            # keep scanning down for the first that fits as fallback
            continue
    if found_fit is None:
        # nothing fits; report the lowest-bpw attempt (least-bad)
        last_q, last_bp = ordered[-1]
        bd = estimate(arch, replace(inp, quant=last_q))
        res.quant = ""
        res.fits = False
        res.note = f"No quant fits the budget at n_ctx={inp.n_ctx}."
        res.fallback_quant = last_q
        res.fallback_bpw = last_bp
        res.fallback_total = bd.total_bytes
        return res
    res.quant, res.bpw, res.total_bytes = found_fit
    res.fits = True
    if fallback is not None:
        res.fallback_quant, res.fallback_bpw, res.fallback_total = fallback
    return res


@dataclass
class MinGpuResult:
    n_gpus: int = 0
    subset: list[float] = field(default_factory=list)  # vram_gb of chosen GPUs
    total_bytes: float = 0.0
    fits: bool = False
    note: str = ""


def min_gpu_setup(arch: ModelArch, inp: Inputs) -> MinGpuResult:
    """Smallest subset (largest-first) of the user's GPUs that fits.

    Greedy: sort the configured GPUs by VRAM descending and add them one at a
    time until the estimate fits. Reports the chosen subset and count.
    """
    if not inp.gpu_vram_gb:
        return MinGpuResult(note="No GPUs configured.")
    # sort largest-first, keep original-size list
    sorted_gpus = sorted(inp.gpu_vram_gb, reverse=True)
    res = MinGpuResult()
    for k in range(1, len(sorted_gpus) + 1):
        subset = sorted_gpus[:k]
        bd = estimate(arch, replace(inp, gpu_vram_gb=subset))
        if _fits(bd):
            return MinGpuResult(
                n_gpus=k, subset=subset, total_bytes=bd.total_bytes,
                fits=True,
                note=f"Fits on {k} GPU(s): {[f'{g}GB' for g in subset]}.",
            )
    # none fit even with all GPUs
    bd = estimate(arch, replace(inp, gpu_vram_gb=sorted_gpus))
    return MinGpuResult(
        n_gpus=len(sorted_gpus), subset=sorted_gpus,
        total_bytes=bd.total_bytes, fits=False,
        note=f"Doesn't fit even on all {len(sorted_gpus)} GPU(s).",
    )