"""Canonical text format shared by SFT data, training and inference.

One turn is `<|user|>QUESTION<|spec|>SPEC<|result|>RESULT<|answer|>ANSWER<|end|>`; a document is
`<|bos|>` + turns + `<|eos|>`. Only SPEC and `ANSWER<|end|>` are trained; RESULT is always
written by the engine.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

from zenith.engines.base import Result
from zenith.engines.registry import schemas
from zenith.tokenizer.bpe import BPETokenizer

IGNORE = -100


def fmt_num(x: float, sig: int = 4) -> str:
    """`sig` significant digits, no trailing zeros; exponent form (1.988e30, 1.23e-4) outside [1e-3, 1e6)."""
    x = float(x)
    if x == 0:
        return "0"
    rounded = float(f"{x:.{sig - 1}e}")
    if abs(rounded) >= 1e6 or abs(rounded) < 1e-3:
        mant, exp = f"{x:.{sig - 1}e}".split("e")
        mant = mant.rstrip("0").rstrip(".") if "." in mant else mant
        return f"{mant}e{int(exp)}"
    digits = max(sig - 1 - math.floor(math.log10(abs(rounded))), 0)
    s = f"{rounded:.{digits}f}"
    return s.rstrip("0").rstrip(".") if "." in s else s


def _value(v) -> str:
    if isinstance(v, str):
        return '"' + v.replace("\\", "\\\\").replace('"', '\\"') + '"'
    return fmt_num(v)


def _obj(pairs) -> str:
    return "{" + ",".join(f'"{k}":{v}' for k, v in pairs) + "}"


def spec_json(spec: dict) -> str:
    """Compact JSON with keys in order sim, op, then the op's schema order."""
    pairs = [("sim", _value(spec["sim"]))]
    if "op" in spec:
        pairs.append(("op", _value(spec["op"])))
        order = [p.name for p in schemas()[spec["sim"]][spec["op"]]]
        pairs += [(k, _value(spec[k])) for k in order if k in spec]
    return _obj(pairs)


def result_json(result: Result) -> str:
    d = result.prompt_dict()
    values = _obj((k, _value(v)) for k, v in d["values"].items())
    notes = "[" + ",".join(_value(n) for n in d["notes"]) + "]"
    return _obj([("values", values), ("notes", notes)])


@dataclass
class Turn:
    question: str
    spec: dict
    result: Result
    answer: str


def _segments(tok: BPETokenizer, turn: Turn) -> list[tuple[list[int], bool]]:
    """(ids, trainable) pieces of one turn."""
    sp = tok.special_ids
    return [
        ([sp["<|user|>"], *tok.encode(turn.question), sp["<|spec|>"]], False),
        (tok.encode(spec_json(turn.spec)), True),
        ([sp["<|result|>"], *tok.encode(result_json(turn.result)), sp["<|answer|>"]], False),
        ([*tok.encode(turn.answer), sp["<|end|>"]], True),
    ]


def encode_turns(tok: BPETokenizer, turns: list[Turn], train_last_only: bool = True) -> tuple[list[int], list[int]]:
    """Token ids and next-token targets; targets are IGNORE except where the next token is trained."""
    ids = [tok.special_ids["<|bos|>"]]
    trained = [False]
    for i, turn in enumerate(turns):
        train_turn = not train_last_only or i == len(turns) - 1
        for seg, trainable in _segments(tok, turn):
            ids += seg
            trained += [trainable and train_turn] * len(seg)
    ids.append(tok.special_ids["<|eos|>"])
    trained.append(False)
    targets = [ids[i + 1] if trained[i + 1] else IGNORE for i in range(len(ids) - 1)] + [IGNORE]
    return ids, targets


def prompt_ids(tok: BPETokenizer, history: list[Turn], question: str) -> list[int]:
    """Prefix for inference: history turns, then the new question, ending at `<|spec|>`."""
    ids = [tok.special_ids["<|bos|>"]]
    for turn in history:
        for seg, _ in _segments(tok, turn):
            ids += seg
    return ids + [tok.special_ids["<|user|>"], *tok.encode(question), tok.special_ids["<|spec|>"]]


def continuation_text(spec_text: str, result: Result) -> str:
    """Text that follows the prompt before answer generation."""
    return f"{spec_text}<|result|>{result_json(result)}<|answer|>"