"""YaRN / RoPE context-extension sanity checks.

YaRN changes how the model *interprets* positions, not the per-token KV size,
so KV cache memory is driven by the target context length the user requests.
These helpers surface warnings when the target context exceeds the model's
training context (YaRN scaling required) and show the effective context as a
sanity check.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class YarnConfig:
    """Recommended YaRN/RoPE parameters for a target context.

    ``rope_freq_scale`` is the only value we actually derive; the rest are the
    llama.cpp defaults (matching :class:`vramcalc.report.Inputs`) so a stale
    manual tweak can't linger after auto-configuration.
    """
    rope_freq_scale: float = 1.0
    yarn_ext_factor: float = -1.0
    yarn_attn_factor: float = 1.0
    yarn_beta_fast: float = 32.0
    yarn_beta_slow: float = 1.0
    scaling: bool = False          # True when target > training_ctx (yarn needed)
    note: str = ""                 # human-readable explanation / caveats


def auto_configure_yarn(training_ctx: int, target_n_ctx: int) -> YarnConfig:
    """Derive recommended YaRN/RoPE params from training_ctx + a target n_ctx.

    The llama.cpp long-context recipe is ``--rope-scaling yarn`` with
    ``rope_freq_scale = training_ctx / target_n_ctx`` (so the effective context
    matches the target), leaving extrapolation/attention/beta at their defaults.
    Returns a :class:`YarnConfig` with everything reset to those defaults when no
    extension is needed or the inputs are unusable.
    """
    if target_n_ctx is None or target_n_ctx <= 0:
        return YarnConfig(note="Target context is invalid; nothing to configure.")
    if training_ctx is None or training_ctx <= 0:
        return YarnConfig(
            note="Training context unknown — fetch a GGUF or load a preset, "
                 "or set rope_freq_scale manually.",
        )
    if target_n_ctx <= training_ctx:
        return YarnConfig(
            note=f"Target context {target_n_ctx} is within training context "
                 f"{training_ctx}; no YaRN scaling needed.",
        )
    scale = round(training_ctx / target_n_ctx, 6)
    if scale <= 0:
        return YarnConfig(note="Target context is invalid; nothing to configure.")
    eff = yarn_effective_context(training_ctx, scale)
    stretch = (1.0 / scale) if scale else 0.0
    return YarnConfig(
        rope_freq_scale=scale,
        scaling=True,
        note=f"Set rope_freq_scale={scale} so effective context ≈ {eff} "
             f"(training_ctx {training_ctx} × {stretch:.1f}×). "
             f"YaRN extrapolation auto (--yarn-ext-factor -1).",
    )


def yarn_effective_context(training_ctx: int, rope_freq_scale: float) -> int:
    """Effective (interpolated) context given a rope_freq_scale.

    effective = training_ctx / rope_freq_scale (rope_freq_scale < 1 extends).
    """
    if rope_freq_scale <= 0:
        return training_ctx
    return int(round(training_ctx / rope_freq_scale))


def yarn_warnings(
    *,
    training_ctx: int,
    target_ctx: int,
    rope_freq_scale: float,
    yarn_ext_factor: float,
    yarn_attn_factor: float,
) -> list[str]:
    """Return human-readable warnings about the YaRN / context config."""
    warns: list[str] = []

    if target_ctx > training_ctx:
        warns.append(
            f"Target context {target_ctx} exceeds training context "
            f"{training_ctx}; YaRN/RoPE scaling is required for coherent "
            f"long-context output."
        )

    eff = yarn_effective_context(training_ctx, rope_freq_scale)
    if target_ctx > training_ctx and eff < target_ctx:
        warns.append(
            f"Effective context from rope_freq_scale={rope_freq_scale} is "
            f"{eff}, which is below the target {target_ctx}. The model may "
            f"not fully cover the requested window; lower rope_freq_scale "
            f"to interpolate more aggressively."
        )

    if target_ctx > training_ctx and yarn_ext_factor < 0.0:
        warns.append("yarn_ext_factor is negative; YaRN extrapolation is off.")

    if yarn_attn_factor < 1.0:
        warns.append(
            "yarn_attn_factor < 1.0 scales attention down; use ~1.0 for "
            "normal long-context YaRN (lower only for high-context YaRN fine-tunes)."
        )

    warns.extend(yarn_coherence_warnings(training_ctx, target_ctx))

    return warns


# Extension-ratio thresholds for the coherence ladder. These are warnings
# only — never a hard block — and come from YaRN-paper behavior and community
# experience, not measurement. The ladder is monotonic: more extension means
# same-or-more concern. See PLAN.md "YaRN coherence ladder".
_YARN_FINE_RATIO = 2.0       # <= 2x: no warning
_YARN_BAD_RATIO = 8.0        # > 2x..8x: tail degradation; >= 8x: incoherent


def extension_ratio(training_ctx: int, target_ctx: int) -> float:
    """How far the target context stretches beyond the training context.

    Returns 1.0 when target <= training (no extension) or when training_ctx
    is unknown/invalid.
    """
    if not training_ctx or training_ctx <= 0:
        return 1.0
    if target_ctx <= 0:
        return 1.0
    if target_ctx <= training_ctx:
        return 1.0
    return target_ctx / training_ctx


def yarn_coherence_warnings(training_ctx: int, target_ctx: int) -> list[str]:
    """Warn about coherence risk when extending context beyond training.

    These are heuristic warnings, not guarantees. The ladder is monotonic:
      <= 2x  : fine, no warning
      > 2x..8x : expect some degradation at the tail of the context window
      >= 8x  : likely incoherent at long range; not recommended
    """
    if training_ctx <= 0 or target_ctx <= 0:
        return []
    if target_ctx <= training_ctx:
        return []
    ratio = target_ctx / training_ctx
    w: list[str] = []
    if ratio >= _YARN_BAD_RATIO:
        w.append(
            f"Target context {target_ctx} is {ratio:.1f}x the training "
            f"context {training_ctx}. Likely incoherent at long range; "
            f"not recommended."
        )
    elif ratio > _YARN_FINE_RATIO:
        w.append(
            f"Target context {target_ctx} is {ratio:.1f}x the training "
            f"context {training_ctx}. Expect some quality degradation at "
            f"the tail of the context window."
        )
    return w