"""
Interactive chat with the exported Smartwatch LM.

Usage:
    pip install torch tokenizers
    python chat.py
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

import torch

import config as cfg
from model import load_model
from reply_utils import build_prompt, extract_bot_reply_from_continuation, extract_intent_reply


class ChatSession:
    def __init__(
        self,
        checkpoint_path: Path | None = None,
        tokenizer_path: Path | None = None,
        device: str | None = None,
        max_new_tokens: int | None = None,
        temperature: float | None = None,
        top_k: int | None = None,
    ):
        self.model, self.tokenizer, self.device = load_model(
            checkpoint_path, tokenizer_path, device
        )
        self.max_new_tokens = max_new_tokens or cfg.SAMPLE_MAX_NEW_TOKENS
        self.temperature = temperature if temperature is not None else cfg.SAMPLE_TEMPERATURE
        self.top_k = top_k if top_k is not None else cfg.SAMPLE_TOP_K
        self.history: list[tuple[str, str]] = []

    def reset(self) -> None:
        self.history.clear()

    @torch.no_grad()
    def say(self, user_message: str) -> str:
        user_message = user_message.strip()
        if not user_message:
            return ""

        prompt = build_prompt(self.history, user_message)
        start_ids = self.tokenizer.encode(prompt).ids
        x = torch.tensor([start_ids], dtype=torch.long, device=self.device)
        y = self.model.generate(
            x,
            max_new_tokens=self.max_new_tokens,
            temperature=self.temperature,
            top_k=self.top_k,
        )
        new_ids = y[0, len(start_ids) :].tolist()
        continuation = self.tokenizer.decode(new_ids)
        reply = extract_bot_reply_from_continuation(continuation)
        self.history.append((user_message, reply))
        return reply

    def say_display(self, user_message: str) -> tuple[str, str, str]:
        """Return (raw_reply, intent, display_text)."""
        raw = self.say(user_message)
        parsed = extract_intent_reply(raw)
        return raw, parsed.intent, parsed.template


def print_banner() -> None:
    print("Smartwatch LM chat — type a message and press Enter.")
    print("Commands: quit/exit  |  reset (clear history)  |  history")
    print("-" * 60)


def run_repl() -> None:
    try:
        session = ChatSession()
    except FileNotFoundError as exc:
        print(exc, file=sys.stderr)
        sys.exit(1)

    val_loss = None
    ckpt_path = cfg.OUTPUT_DIR / "checkpoint.pt"
    if ckpt_path.is_file():
        checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        val_loss = checkpoint.get("best_val_loss")

    print_banner()
    print(f"device: {session.device}")
    if val_loss is not None:
        print(f"checkpoint val loss: {val_loss:.4f}")
    print()

    while True:
        try:
            user_input = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nbye")
            break

        if not user_input:
            continue
        lowered = user_input.lower()
        if lowered in {"quit", "exit"}:
            print("bye")
            break
        if lowered == "reset":
            session.reset()
            print("(history cleared)")
            continue
        if lowered == "history":
            if not session.history:
                print("(empty)")
            for user_text, bot_text in session.history:
                print(f"user: {user_text}\nbot: {bot_text}\n")
            continue

        _, intent, display = session.say_display(user_input)
        print(f"bot> {display}")
        if intent and intent != "NONE":
            print(f"     intent: {intent}")


if __name__ == "__main__":
    run_repl()