import copy
import json
import os
from typing import Dict, Any

import torch

from config import CONFIG
from inference import _build_tokenizers, _resolve_device, load_model, run_inference


_STATE = {
    "loaded": False,
    "model": None,
    "cfg": None,
    "device": None,
    "src_tok": None,
    "tgt_tok": None,
}


def _read_model_settings() -> Dict[str, Any]:
    if not os.path.exists("model_settings.json"):
        return {}
    try:
        with open("model_settings.json", "r", encoding="utf-8") as f:
            data = json.load(f)
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def _load_once() -> None:
    if _STATE["loaded"]:
        return

    settings = _read_model_settings()
    cfg = copy.deepcopy(CONFIG)
    cfg["model_type"] = os.environ.get(
        "HF_MODEL_TYPE",
        settings.get("model_type", "d3pm_cross_attention"),
    )
    cfg["data"]["include_negative_examples"] = (
        os.environ.get(
            "HF_INCLUDE_NEG",
            str(settings.get("include_negative_examples", True)).lower(),
        ).lower()
        == "true"
    )
    num_steps_raw = os.environ.get("HF_NUM_STEPS", settings.get("num_steps"))
    if num_steps_raw is not None:
        num_steps = int(num_steps_raw)
        cfg["model"]["diffusion_steps"] = num_steps
        cfg["inference"]["num_steps"] = num_steps
    device = _resolve_device(cfg)

    model, cfg = load_model("best_model.pt", cfg, device)
    src_tok, tgt_tok = _build_tokenizers(cfg)

    _STATE["model"] = model
    _STATE["cfg"] = cfg
    _STATE["device"] = device
    _STATE["src_tok"] = src_tok
    _STATE["tgt_tok"] = tgt_tok
    _STATE["loaded"] = True


def _clean_text(text: str) -> str:
    text = " ".join(text.split())
    if not text:
        return text
    toks = text.split()
    out = []
    prev = None
    run = 0
    for tok in toks:
        if tok == prev:
            run += 1
        else:
            prev = tok
            run = 1
        if run <= 2:
            out.append(tok)
    s = " ".join(out)
    s = s.replace(" ।", "।").replace(" ॥", "॥")
    return " ".join(s.split())


def predict(
    text: str,
    temperature: float = 0.7,
    top_k: int = 40,
    repetition_penalty: float = 1.2,
    diversity_penalty: float = 0.0,
    num_steps: int = 64,
    clean_output: bool = True,
) -> Dict[str, Any]:
    _load_once()
    if not text or not text.strip():
        return {"error": "empty input", "output": ""}

    cfg = copy.deepcopy(_STATE["cfg"])
    cfg["inference"]["temperature"] = float(temperature)
    cfg["inference"]["top_k"] = int(top_k)
    cfg["inference"]["repetition_penalty"] = float(repetition_penalty)
    cfg["inference"]["diversity_penalty"] = float(diversity_penalty)
    cfg["inference"]["num_steps"] = int(num_steps)

    src_tok = _STATE["src_tok"]
    tgt_tok = _STATE["tgt_tok"]
    device = _STATE["device"]

    input_ids = torch.tensor([src_tok.encode(text.strip())], dtype=torch.long, device=device)
    out = run_inference(_STATE["model"], input_ids, cfg)
    decoded_ids = [x for x in out[0].tolist() if x > 4]
    raw = tgt_tok.decode(decoded_ids).strip()
    output = _clean_text(raw) if clean_output else raw

    return {
        "input": text,
        "output": output,
        "raw_output": raw,
        "config": {
            "temperature": float(temperature),
            "top_k": int(top_k),
            "repetition_penalty": float(repetition_penalty),
            "diversity_penalty": float(diversity_penalty),
            "num_steps": int(num_steps),
            "clean_output": bool(clean_output),
        },
    }