"""Shared utilities: device, seed, prompt format, reward, checkpoint load/save."""
from __future__ import annotations

import os
import random
import re
from contextlib import nullcontext

import numpy as np
import torch

from tokenizer import ASSISTANT, EOS, USER

# ------------------------------------------------------------------ runtime

def get_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def autocast_ctx(device: str):
    """Mixed precision: bf16 where the GPU supports it (A100/L4), fp16 on T4, fp32 on CPU."""
    if device != "cuda":
        return nullcontext(), None
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype), dtype


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def ensure_dir(path: str) -> None:
    d = os.path.dirname(path)
    if d:
        os.makedirs(d, exist_ok=True)


# ------------------------------------------------------------ chat format
PROMPT_TEMPLATE = "Write a short story that uses the words: {words}."


def build_prompt(words: list[str]) -> str:
    return f"{USER}{PROMPT_TEMPLATE.format(words=', '.join(words))}{ASSISTANT}"


def build_example(words: list[str], story: str) -> tuple[str, str]:
    """Return (prompt, response). The SFT loss is computed on the response only."""
    return build_prompt(words), story.strip() + EOS


_STOP = set("""
a an the and or but so if then than that this these those there their they them he she his her him it its
was were is are be been being am do does did done have has had having to of in on at for with from by as
into onto over under up down out off very too also just not no yes all any some one two three every each
said says say went go goes going came come comes get got make made can could would should will shall may
might must when where what which who whom why how because while after before again once day time little
big lot lots very much many more most other another then now away back here there with
upon named lived felt wanted liked loved looked found know knew think thought told asked started saw
""".split())


def pick_keywords(story: str, k: int, rng: random.Random) -> list[str] | None:
    """Pick k content words from a story — used as the SFT/DPO instruction."""
    cands = sorted({w for w in re.findall(r"[a-z]+", story.lower()) if len(w) >= 4 and w not in _STOP})
    if len(cands) < k:
        return None
    return rng.sample(cands, k)


# ----------------------------------------------------------------- reward
def keyword_reward(text: str, words: list[str], finished: bool) -> float:
    """
    Rule-based, verifiable reward (in the spirit of 'verifiable rewards'):
      - fraction of required words that appear (0..1)
      - +0.2 if the model closes the story with EOS (not truncated)
    """
    found = set(re.findall(r"[a-z]+", text.lower()))
    hit = sum(1 for w in words if w in found) / max(len(words), 1)
    return hit + (0.2 if finished else 0.0)


def keyword_hit_rate(text: str, words: list[str]) -> float:
    found = set(re.findall(r"[a-z]+", text.lower()))
    return sum(1 for w in words if w in found) / max(len(words), 1)


def keyword_mentions(text: str, words: list[str]) -> float:
    """Average number of times each required word appears. ~1-2 is natural; much higher = stuffing."""
    toks = re.findall(r"[a-z]+", text.lower())
    return sum(toks.count(w) for w in words) / max(len(words), 1)


def repeated_trigram_rate(text: str) -> float:
    """Fraction of word trigrams that already appeared earlier in the text (degenerate repetition)."""
    toks = re.findall(r"[a-z']+", text.lower())
    grams = list(zip(toks, toks[1:], toks[2:]))
    if not grams:
        return 0.0
    return 1.0 - len(set(grams)) / len(grams)


def keyword_reward_v2(text: str, words: list[str], finished: bool) -> float:
    """
    Anti-stuffing reward: v1 plus penalties for the failure modes seen in the first run.
      - penalty 0.1 per mention beyond 2 per required word (keyword stuffing)
      - penalty 1.0 x repeated-trigram rate (looping text)
    """
    toks = re.findall(r"[a-z]+", text.lower())
    excess = sum(max(0, toks.count(w) - 2) for w in words)
    return keyword_reward(text, words, finished) - 0.1 * excess - repeated_trigram_rate(text)


REWARDS = {"v1": keyword_reward, "v2": keyword_reward_v2}


# ------------------------------------------------------------ checkpoints
def save_checkpoint(path: str, model, model_cfg, extra: dict | None = None) -> None:
    ensure_dir(path)
    raw = getattr(model, "_orig_mod", model)  # unwrap torch.compile
    torch.save({"model": raw.state_dict(), "model_cfg": model_cfg.__dict__, **(extra or {})}, path)


def load_model(path: str, device: str):
    from config import ModelConfig
    from model import GPT

    ckpt = torch.load(path, map_location=device, weights_only=False)
    cfg = ModelConfig(**ckpt["model_cfg"])
    model = GPT(cfg).to(device)
    model.load_state_dict(ckpt["model"])
    return model, cfg


def load_model_dir(path: str, device: str):
    """Load an exported model folder (config.json + model.safetensors), e.g. from the Hugging Face Hub."""
    import json

    from safetensors.torch import load_model as st_load

    from config import ModelConfig
    from model import GPT

    with open(os.path.join(path, "config.json")) as f:
        cfg = ModelConfig(**json.load(f))
    model = GPT(cfg)
    st_load(model, os.path.join(path, "model.safetensors"))
    return model.to(device).eval(), cfg