"""CBT-18.95M: character-level Transformer with Decoupled 2D-RoPE.

Quick start:
    from cbt import load, add, chat
    model, tok = load()                      # local folder or the Hugging Face repo
    add(model, tok, 123456789, 987654321)    # {'answer': '1111111110', 'correct': True, ...}
    chat(model, tok, "What is the weather in Tokyo?")
    twin, _ = load(variant="twin_1d")        # matched 1D-RoPE control (same architecture, data, steps)
"""

import json
import os
from dataclasses import fields
from typing import Any, Callable, Dict, Optional, Tuple, Union

import torch

from .model import ScaledGPT as CBT, ScaledGPTConfig as CBTConfig
from .tokenizer import ScaledTokenizer as CBTTokenizer
from .tools import execute_tool

__all__ = ["CBT", "CBTConfig", "CBTTokenizer", "load", "add", "chat", "execute_tool"]

REPO_ID = "Neithanvarvas/CBT"
VARIANTS: Dict[str, str] = {"cbt_2d": "", "twin_1d": "twin_1d/"}
_PKG_DIR = os.path.dirname(os.path.abspath(__file__))


def _resolve(source: Optional[str], filename: str) -> str:
    """Finds `filename` in a local folder (default: the repo root next to this package) or downloads it."""
    folder = source or os.path.dirname(_PKG_DIR)
    local = os.path.join(folder, filename)
    # A clone made without git-lfs leaves a ~130-byte pointer file instead of the weights.
    if os.path.isfile(local) and not (filename.endswith(".safetensors") and os.path.getsize(local) < 1024):
        return local
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id=source if source and not os.path.isdir(source) else REPO_ID, filename=filename)


def load(
    source: Optional[str] = None,
    device: Optional[Union[str, torch.device]] = None,
    variant: str = "cbt_2d",
) -> Tuple[CBT, CBTTokenizer]:
    """Loads model and tokenizer from a local folder or a Hugging Face repo id; returns (model, tokenizer) in eval mode.

    variant: "cbt_2d" (the released model) or "twin_1d" (its 1D-RoPE control).
    """
    from safetensors.torch import load_file

    if variant not in VARIANTS:
        raise ValueError(f"variant must be one of {sorted(VARIANTS)}")
    prefix = VARIANTS[variant]
    with open(_resolve(source, prefix + "config.json"), encoding="utf-8") as f:
        raw: Dict[str, Any] = json.load(f)
    config = CBTConfig(**{k.name: raw[k.name] for k in fields(CBTConfig) if k.name in raw})

    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    tok = CBTTokenizer()
    model = CBT(config, tokenizer=tok)
    model.load_state_dict(load_file(_resolve(source, prefix + "model.safetensors")))
    return model.to(device).eval(), tok


@torch.no_grad()
def _generate(model: CBT, ids: torch.Tensor, max_new_tokens: int, stop: list) -> torch.Tensor:
    return model.generate(ids, max_new_tokens=max_new_tokens, temperature=0.0, stop_tokens=stop)


def _encode(model: CBT, tok: CBTTokenizer, text: str, bos: bool = True) -> torch.Tensor:
    device = next(model.parameters()).device
    return torch.tensor([tok.encode(text, add_bos=bos, add_eos=False)], dtype=torch.long, device=device)


def add(model: CBT, tok: CBTTokenizer, a: Union[int, str], b: Union[int, str]) -> Dict[str, Any]:
    """Greedy-decodes a + b with the column scratchpad. Trained on 1-25 digit operands."""
    a, b = str(a).strip(), str(b).strip()
    if not (a.isdigit() and b.isdigit()):
        raise ValueError("operands must be non-negative integers")
    budget = 12 * (max(len(a), len(b)) + 2) + 8
    out = _generate(model, _encode(model, tok, f"{a}+{b}=<THINK>"), budget, [tok.eos_id])
    raw = tok.decode(out[0].tolist())
    think, _, tail = raw.partition("<THINK>")[2].partition("</THINK>")
    answer = tail.lstrip("=").replace("<EOS>", "").strip()
    return {"answer": answer, "scratchpad": think, "correct": answer == str(int(a) + int(b)), "raw": raw}


def chat(
    model: CBT,
    tok: CBTTokenizer,
    message: str,
    tool_executor: Optional[Callable[[str], str]] = execute_tool,
    max_new_tokens: int = 256,
) -> Dict[str, Any]:
    """Runs one user turn: <THINK> reasoning, optional <CALL> executed by `tool_executor`, then <ANSWER>."""
    ids = _encode(model, tok, f"<USER>{message}</USER><THINK>")
    out = _generate(model, ids, max_new_tokens, [tok.call_e_id, tok.answer_e_id, tok.eos_id])
    raw = tok.decode(out[0].tolist())

    if "</CALL>" in raw and tool_executor is not None:
        call = raw.split("<CALL>")[-1].split("</CALL>")[0]
        cont = torch.cat([out, _encode(model, tok, f"<RESULT>{tool_executor(call)}</RESULT>", bos=False)], dim=1)
        out = _generate(model, cont, max_new_tokens, [tok.answer_e_id, tok.eos_id])
        raw = tok.decode(out[0].tolist())

    def span(tag: str) -> Optional[str]:
        s, e = f"<{tag}>", f"</{tag}>"
        return raw.split(s, 1)[1].split(e, 1)[0].strip() if s in raw and e in raw else None

    return {"think": span("THINK"), "call": span("CALL"), "result": span("RESULT"), "answer": span("ANSWER"), "raw": raw}