"""Inference engine (PRIVATE to Serving_Module).

* Modes: non_thinking (direct answer), thinking (<think>...</think> then answer), max_thinking (long reasoning).
  The mode is selected through the same `[MODE: ...]` system-prompt tag the model was fine-tuned with.
* Lag-free: KV-cache generation in a worker thread + TextIteratorStreamer, so tokens appear as they are produced.
* Stops cleanly on <|im_end|> (the stop token the SFT stage aligned) and never prints special tokens.
* Live `teach`: a question/answer pair is used IMMEDIATELY (direct hit or injected as a known fact) and is also
  appended to chat.jsonl in the exact SFT schema, so the next training round learns it permanently.

The chat format is intentionally duplicated here (it mirrors Training_Module's tokenizer_surgeon) because
modules may not import each other's internals.
"""
from __future__ import annotations

import difflib
import json
import re
import threading
import time
from collections import deque
from pathlib import Path
from typing import Callable, Deque, Dict, List, Optional, Tuple

IM_START, IM_END, ENDOFTEXT = "<|im_start|>", "<|im_end|>", "<|endoftext|>"
THINK_OPEN, THINK_CLOSE = "<think>", "</think>"
MODES = ("thinking", "non_thinking", "max_thinking")
DEFAULT_SYSTEM = "You are RSS AI, a helpful assistant created by RSS. You understand Hindi, Hinglish and English."
TOKEN_BUDGET = {"non_thinking": 220, "thinking": 420, "max_thinking": 800}
SAMPLING = {"non_thinking": 0.45, "thinking": 0.6, "max_thinking": 0.65}
_WS = re.compile(r"\s+")
_SPECIAL = re.compile(re.escape(IM_END) + "|" + re.escape(ENDOFTEXT) + "|" + re.escape(IM_START))


def normalise_mode(value: str) -> Optional[str]:
    value = (value or "").strip().lower().replace("-", "_").replace(" ", "_").lstrip("/")
    return value if value in MODES else None


def _norm(text: str) -> str:
    return _WS.sub(" ", text.strip()).casefold()


def split_thinking(text: str) -> Tuple[str, str, bool]:
    """-> (thinking, answer, truncated). `truncated` = the token budget ended inside <think>."""
    text = _SPECIAL.sub("", text).strip()
    if THINK_CLOSE in text:
        head, _, answer = text.partition(THINK_CLOSE)
        return head.replace(THINK_OPEN, "").strip(), answer.strip(), False
    if THINK_OPEN in text:
        return text.replace(THINK_OPEN, "").strip(), "", True
    return "", text, False


class InferenceEngine:
    def __init__(self, model, tokenizer, chat_file: Optional[str] = None, max_new_tokens: Optional[int] = None,
                 history_turns: int = 3, active_vocab_file: Optional[str] = None):
        import torch
        self.torch = torch
        self.model, self.tokenizer = model, tokenizer
        self.model.eval()
        try:
            self.model.config.use_cache = True
        except Exception:
            pass
        try:
            self.device = next(model.parameters()).device
        except StopIteration:
            self.device = torch.device("cpu")
        vocab = tokenizer.get_vocab()
        self.eos_ids = [vocab[t] for t in (IM_END, ENDOFTEXT) if t in vocab] or [tokenizer.eos_token_id]
        self.pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else self.eos_ids[0]
        self.max_new_tokens = max_new_tokens
        self.active_token_ids: Optional[List[int]] = None
        if active_vocab_file:
            vocab_path = Path(active_vocab_file)
            if vocab_path.is_file():
                try:
                    self.active_token_ids = sorted({int(x) for x in json.loads(vocab_path.read_text(encoding="utf-8"))})
                    print(f"[Serving] active vocabulary mask loaded: {len(self.active_token_ids):,} token IDs", flush=True)
                except Exception as exc:
                    raise RuntimeError(f"could not load active vocabulary manifest {vocab_path}: {exc}") from exc
        self.chat_file = Path(chat_file) if chat_file else None
        self.history: Dict[str, Deque[Tuple[str, str]]] = {}
        self.history_turns = history_turns
        self.taught: Dict[str, Tuple[str, str]] = {}
        self.lock = threading.Lock()                      # one generation at a time (CPU / shared weights)
        self._load_taught()

    # ------------------------------------------------------------------ teach
    def _load_taught(self) -> None:
        if not self.chat_file or not self.chat_file.exists():
            return
        try:
            with open(self.chat_file, "r", encoding="utf-8") as handle:
                for line in handle:
                    if '"teach-' not in line:
                        continue
                    try:
                        rec = json.loads(line)
                    except ValueError:
                        continue
                    if str(rec.get("id", "")).startswith("teach-") and rec.get("user") and rec.get("assistant"):
                        self.taught[_norm(rec["user"])] = (rec["user"], rec["assistant"])
        except OSError:
            pass

    def teach(self, question: str, answer: str) -> str:
        question, answer = question.strip(), answer.strip()
        if not question or not answer:
            return "teach: needs both a question and an answer ->  teach: सवाल | जवाब"
        self.taught[_norm(question)] = (question, answer)
        saved = ""
        if self.chat_file:
            record = {"id": f"teach-{int(time.time() * 1000)}", "mode": "non_thinking", "system": DEFAULT_SYSTEM,
                      "user": question, "thinking": "", "assistant": answer}
            try:
                self.chat_file.parent.mkdir(parents=True, exist_ok=True)
                with open(self.chat_file, "a", encoding="utf-8") as handle:
                    handle.write(json.dumps(record, ensure_ascii=False) + "\n")
                saved = f" and saved to {self.chat_file.name} (next training round will learn it)"
            except OSError as exc:
                saved = f" (could not write {self.chat_file.name}: {exc})"
        return f"learned: {question!r} -> {answer!r}{saved}"

    def _taught_lookup(self, question: str) -> Tuple[Optional[str], List[Tuple[str, str]]]:
        key = _norm(question)
        if key in self.taught:
            return self.taught[key][1], []
        scored = sorted(((difflib.SequenceMatcher(None, key, k).ratio(), k) for k in self.taught), reverse=True)
        if scored and scored[0][0] >= 0.90:
            return self.taught[scored[0][1]][1], []
        related = [self.taught[k] for ratio, k in scored[:3] if ratio >= 0.55]
        return None, related

    # ------------------------------------------------------------------ prompt
    def _render(self, user: str, mode: str, session: str, facts: List[Tuple[str, str]]) -> str:
        system = DEFAULT_SYSTEM
        if facts:
            system += "\nKnown facts taught by the user: " + " | ".join(f"{q} -> {a}" for q, a in facts)
        parts = [f"{IM_START}system\n{system}\n[MODE: {mode}]{IM_END}\n"]
        for past_user, past_answer in self.history.get(session, []):
            parts.append(f"{IM_START}user\n{past_user}{IM_END}\n{IM_START}assistant\n{past_answer}{IM_END}\n")
        parts.append(f"{IM_START}user\n{user}{IM_END}\n{IM_START}assistant\n")
        return "".join(parts)

    def reset(self, session: str = "default") -> None:
        self.history.pop(session, None)

    # ------------------------------------------------------------------ generation
    def generate(self, user: str, mode: str = "non_thinking", session: str = "default",
                 on_text: Optional[Callable[[str], None]] = None) -> Dict:
        mode = normalise_mode(mode) or "non_thinking"
        started = time.time()
        direct, facts = self._taught_lookup(user)
        if direct is not None:
            if on_text:
                on_text(direct)
            self._remember(session, user, direct)
            return {"mode": mode, "thinking": "", "answer": direct, "truncated": False, "source": "taught",
                    "tokens": 0, "seconds": round(time.time() - started, 2)}
        prompt = self._render(user, mode, session, facts)
        budget = self.max_new_tokens or TOKEN_BUDGET[mode]
        with self.lock:
            raw, tokens = self._run(prompt, mode, budget, on_text)
        thinking, answer, truncated = split_thinking(raw)
        if truncated or (not answer and mode != "non_thinking"):
            answer = answer or "(token limit reached while thinking - ask again, or use mode non_thinking for a direct answer)"
        self._remember(session, user, answer)
        return {"mode": mode, "thinking": thinking, "answer": answer, "truncated": truncated, "source": "model",
                "tokens": tokens, "seconds": round(time.time() - started, 2)}

    def _remember(self, session: str, user: str, answer: str) -> None:
        turns = self.history.setdefault(session, deque(maxlen=self.history_turns))
        turns.append((user, answer))

    def _run(self, prompt: str, mode: str, budget: int, on_text) -> Tuple[str, int]:
        torch = self.torch
        enc = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(self.device)
        kwargs = dict(**enc, max_new_tokens=budget, do_sample=True, temperature=SAMPLING[mode], top_p=0.9,
                      repetition_penalty=1.08, eos_token_id=self.eos_ids, pad_token_id=self.pad_id, use_cache=True)
        if self.active_token_ids:
            try:
                from transformers import LogitsProcessor, LogitsProcessorList
            except ImportError as exc:
                raise RuntimeError("transformers is required to enforce the active-vocabulary generation mask") from exc
            allowed_ids = self.active_token_ids

            class _ActiveVocabLogitsProcessor(LogitsProcessor):
                def __init__(self, ids):
                    self.ids = ids
                    self._masks = {}

                def __call__(self, input_ids, scores):
                    key = (str(scores.device), int(scores.shape[-1]))
                    mask = self._masks.get(key)
                    if mask is None:
                        if self.ids[-1] >= scores.shape[-1]:
                            raise RuntimeError("active vocabulary manifest contains an ID outside model logits")
                        mask = torch.zeros(scores.shape[-1], dtype=torch.bool, device=scores.device)
                        mask[torch.as_tensor(self.ids, dtype=torch.long, device=scores.device)] = True
                        self._masks[key] = mask
                    scores.masked_fill_(~mask.unsqueeze(0), -float("inf"))
                    return scores

            kwargs["logits_processor"] = LogitsProcessorList([_ActiveVocabLogitsProcessor(allowed_ids)])
        if on_text is None:
            with torch.inference_mode():
                out = self.model.generate(**kwargs)
            new_ids = out[0][enc["input_ids"].shape[1]:]
            return self.tokenizer.decode(new_ids, skip_special_tokens=False), int(new_ids.shape[0])

        from transformers import TextIteratorStreamer
        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=False, timeout=900.0)
        failure: List[BaseException] = []

        def worker() -> None:
            try:
                with torch.inference_mode():
                    self.model.generate(**kwargs, streamer=streamer)
            except BaseException as exc:                   # unblock the consumer, then re-raise below
                failure.append(exc)
                streamer.text_queue.put(streamer.stop_signal)

        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        collected: List[str] = []
        for chunk in streamer:
            collected.append(chunk)
            shown = _SPECIAL.sub("", chunk)
            if shown:
                on_text(shown)
        thread.join()
        if failure:
            raise failure[0]
        raw = "".join(collected)
        return raw, len(self.tokenizer(raw, add_special_tokens=False)["input_ids"])