import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

import torch
from ml.model import TheAbyssModel, AbyssConfig
from ml.tokenizer.tokenizer import AbyssTokenizer
from ml.tokenizer.special_tokens import DEFAULT_SYSTEM


_model = None
_tok = None
_device = None
_cfg = None


def load_model(checkpoint_path: str) -> None:
    global _model, _tok, _device, _cfg
    _device = "cuda" if torch.cuda.is_available() else "cpu"
    ckpt = torch.load(checkpoint_path, map_location=_device, weights_only=False)
    cfg_raw = ckpt["config"]
    _cfg = AbyssConfig(**cfg_raw) if isinstance(cfg_raw, dict) else cfg_raw
    _model = TheAbyssModel(_cfg)
    _model.load_state_dict(ckpt["model"])
    _model.to(_device)
    _model.eval()
    tok_path = os.path.join(os.path.dirname(__file__), "..", "..", "ml", "tokenizer", "vocab")
    _tok = AbyssTokenizer.load(tok_path)
    print(f"[inference] model loaded on {_device}")


def _build_prompt(messages: list) -> str:
    # The model was SFT'd with a system turn on EVERY example. If the client sent
    # no system message, inject the default one so the served prompt matches the
    # training distribution (otherwise the model drifts off the user's request).
    if not any(m.get("role") == "system" for m in messages):
        messages = [{"role": "system", "content": DEFAULT_SYSTEM}] + list(messages)

    parts = []
    for m in messages:
        role = m["role"]
        content = m["content"]
        if role == "system":
            parts.append(f"<|system|>{content}<|endofturn|>")
        elif role == "user":
            parts.append(f"<|user|>{content}<|endofturn|>")
        else:
            parts.append(f"<|assistant|>{content}<|endofturn|>")
    parts.append("<|assistant|>")
    return "".join(parts)


def generate_tokens(
    messages: list,
    max_tokens: int = 200,
    temperature: float = 0.7,
    top_k: int = 40,
    top_p: float = 0.95,
    repetition_penalty: float = 1.2,
):
    """Sync generator - yields decoded token strings one at a time."""
    prompt = _build_prompt(messages)
    ids = _tok.encode(prompt)
    block_size = _cfg.block_size
    # Never ask for more output than the context can hold, and always leave
    # room for the prompt so the input can't collapse to an empty sequence.
    max_tokens = max(1, min(max_tokens, block_size - 8))
    keep = block_size - max_tokens
    if len(ids) > keep:
        ids = ids[-keep:]

    idx = torch.tensor([ids], dtype=torch.long, device=_device)

    eot_ids = _tok.encode("<|endofturn|>")
    eos_ids = _tok.encode("<|endoftext|>")
    stop_ids = set()
    if eot_ids:
        stop_ids.add(eot_ids[0])
    if eos_ids:
        stop_ids.add(eos_ids[0])

    generated: list[int] = []   # tokens produced so far (for repetition penalty)

    with torch.no_grad():
        for _ in range(max_tokens):
            logits, _ = _model(idx[:, -block_size:])
            logits = logits[0, -1, :] / max(temperature, 1e-8)

            # Repetition penalty (CTRL-style) over already-generated tokens only,
            # so we damp loops without suppressing words from the user's prompt.
            if repetition_penalty and repetition_penalty != 1.0 and generated:
                for tid in set(generated):
                    if logits[tid] > 0:
                        logits[tid] /= repetition_penalty
                    else:
                        logits[tid] *= repetition_penalty

            if top_k > 0:
                k = min(top_k, logits.size(-1))
                topk_vals, _ = torch.topk(logits, k)
                logits = logits.masked_fill(logits < topk_vals[-1], float("-inf"))

            probs = torch.softmax(logits, dim=-1)
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cum = torch.cumsum(sorted_probs, dim=-1)
            remove = (cum - sorted_probs) > top_p
            sorted_probs[remove] = 0.0
            sorted_probs = sorted_probs / sorted_probs.sum()

            next_rel = torch.multinomial(sorted_probs, 1)
            next_tok = sorted_idx[next_rel]
            tok_id = next_tok.item()

            if tok_id in stop_ids:
                break

            generated.append(tok_id)
            idx = torch.cat([idx, next_tok.view(1, 1)], dim=1)
            yield _tok.decode([tok_id])