#!/usr/bin/env python3
"""Run Aurora Proelia ChatML locally or start an interactive chat."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from safetensors.torch import load_file
from tokenizers import Tokenizer

from aurora.config import load_model_config
from aurora.model import AuroraForCausalLM


DEFAULT_SYSTEM = (
    "You are Ember Proelia, a proprietary language model created by North ML. "
    "Answer directly and concisely. Do not claim web access or certainty you do not have."
)


def render_chat(messages: list[dict[str, str]], add_generation_prompt: bool = True) -> str:
    text = "".join(
        f"<|im_start|>{item['role']}\n{item['content'].strip()}<|im_end|>\n"
        for item in messages
    )
    if add_generation_prompt:
        text += "<|im_start|>assistant\n"
    return text


def choose_device(value: str) -> torch.device:
    if value != "auto":
        return torch.device(value)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def load_model(root: Path, device: torch.device):
    config = load_model_config(root / "model_ember_proelia_207m_16k.yaml")
    tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
    dtype = torch.float16 if device.type in {"cuda", "mps"} else torch.float32
    model = AuroraForCausalLM(config).to(device=device, dtype=dtype).eval()
    state = load_file(str(root / "model.safetensors"), device=str(device))
    missing, unexpected = model.load_state_dict(state, strict=False)
    missing = [name for name in missing if not name.endswith("._extra_state")]
    if missing or unexpected:
        raise RuntimeError(f"checkpoint mismatch: missing={missing}, unexpected={unexpected}")
    return model, tokenizer, config


def generate(model, tokenizer: Tokenizer, config, prompt: str, device: torch.device, max_new_tokens: int) -> str:
    bos = tokenizer.token_to_id("<bos>")
    eos = tokenizer.token_to_id("<eos>")
    ids = [bos, *tokenizer.encode(prompt, add_special_tokens=False).ids]
    generated: list[int] = []
    with torch.inference_mode():
        for _ in range(max_new_tokens):
            inputs = torch.tensor([ids[-int(config.context_length):]], dtype=torch.long, device=device)
            logits, _ = model(inputs)
            token = int(torch.argmax(logits[0, -1]).item())
            if token == eos:
                break
            generated.append(token)
            ids.append(token)
            text = tokenizer.decode(generated, skip_special_tokens=True)
            if "<|im_end|>" in text or "<|im_start|>" in text:
                break
    text = tokenizer.decode(generated, skip_special_tokens=True).strip()
    for marker in ("<|im_end|>", "<|im_start|>"):
        if marker in text:
            text = text.split(marker, 1)[0].strip()
    return text


def main() -> None:
    parser = argparse.ArgumentParser(description="Aurora Proelia ChatML inference")
    parser.add_argument("--prompt", help="one prompt; omit for interactive chat")
    parser.add_argument("--system", default=DEFAULT_SYSTEM)
    parser.add_argument("--checkpoint-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda", "mps"))
    parser.add_argument("--max-new-tokens", type=int, default=96)
    args = parser.parse_args()
    device = choose_device(args.device)
    model, tokenizer, config = load_model(args.checkpoint_dir, device)
    history: list[dict[str, str]] = [{"role": "system", "content": args.system}]

    def answer(user_text: str) -> str:
        history.append({"role": "user", "content": user_text})
        prompt = render_chat(history)
        response = generate(model, tokenizer, config, prompt, device, args.max_new_tokens)
        history.append({"role": "assistant", "content": response})
        return response

    if args.prompt:
        print(answer(args.prompt))
        return
    print(f"Aurora Proelia ChatML · device={device}")
    print("Type /quit to exit, /clear to reset the conversation.")
    while True:
        try:
            user_text = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if user_text == "/quit":
            break
        if user_text == "/clear":
            history[:] = [{"role": "system", "content": args.system}]
            print("Conversation cleared.")
            continue
        if user_text:
            print(f"Aurora: {answer(user_text)}")


if __name__ == "__main__":
    main()