"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.

Strategy: iterate over candidate (CP, TP, PP) triples in order of
increasing total PE count and return the first one that satisfies
memory + kernel-support constraints with >10% slack. If none fits, return
the best-effort configuration and flag over-budget.
"""
from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Iterable

from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
from .memory_layout import compute_memory


# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
_PP_OPTIONS = (1, 2, 4, 8, 16, 32)


@dataclass
class Suggestion:
    cp: int
    tp: int
    pp: int
    weights_gb: float
    kv_gb: float
    transient_gb: float
    slack_gb: float
    fits: bool
    pes_used: int
    sips_used: int
    reason: str = ""


def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
    """Yield (CP, TP, PP) in order of increasing PE count."""
    triples: list[tuple[int, int, int]] = []
    for tp in _TP_OPTIONS:
        for cp in _CP_OPTIONS:
            for pp in _PP_OPTIONS:
                # PP must not exceed layer count.
                if pp > model.layers:
                    continue
                # Skip TP > H_q * some factor (unrealistic).
                if tp > model.h_q * 4:
                    continue
                triples.append((cp, tp, pp))
    # Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
    # (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
    triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
    for t in triples:
        yield t


def _score_candidate(cp: int, tp: int, pp: int,
                     model: ModelConfig, machine: MachineParams,
                     s_kv: int, mode: str,
                     slack_frac: float = 0.10) -> Suggestion:
    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode)
    cfg = FullConfig(model=model, topo=topo, machine=machine)
    mem = compute_memory(cfg)

    fits = (not mem.over_budget
            and mem.slack_bytes >= slack_frac * mem.budget_bytes)
    reason = ""
    if mem.over_budget:
        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
                  f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
    elif not fits:
        reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
    return Suggestion(
        cp=cp, tp=tp, pp=pp,
        weights_gb=mem.weights_bytes / 1e9,
        kv_gb=mem.kv_cache_bytes / 1e9,
        transient_gb=mem.transient_bytes / 1e9,
        slack_gb=mem.slack_bytes / 1e9,
        fits=fits,
        pes_used=topo.total_pes,
        sips_used=topo.sips_used,
        reason=reason,
    )


def auto_suggest(model: ModelConfig, machine: MachineParams,
                 s_kv: int, mode: str = "decode",
                 slack_frac: float = 0.10) -> Suggestion:
    """Return the smallest-PE-count (CP, TP, PP) that fits.

    If no candidate fits, returns the best-effort one (highest slack,
    even if negative) with fits=False.
    """
    best_fit: Suggestion | None = None
    best_effort: Suggestion | None = None

    for cp, tp, pp in _iter_candidates(model):
        s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac)
        if s.fits and best_fit is None:
            best_fit = s
            break  # candidates are pre-sorted by PE count
        if best_effort is None or s.slack_gb > best_effort.slack_gb:
            best_effort = s

    return best_fit or best_effort  # type: ignore