import re
import time
import torch

class SynapseController:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.device = next(model.parameters()).device
        self.eot_id = tokenizer.convert_tokens_to_ids("<|eot|>")

    def _prompt(self, question):
        return f"<|user|>\n{question}\n<|assistant|>\n"

    def _clean(self, text):
        text = str(text)
        if "<|eot|>" in text:
            text = text.split("<|eot|>")[0]
        if "<|final|>" in text:
            text = text.split("<|final|>")[-1]
        return text.strip()

    def _model(self, question, max_new_tokens=140):
        inputs = self.tokenizer(self._prompt(question), return_tensors="pt").to(self.device)
        input_len = inputs["input_ids"].shape[-1]
        start = time.time()
        with torch.no_grad():
            out = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                min_new_tokens=4,
                do_sample=False,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.eot_id,
                repetition_penalty=1.12,
                no_repeat_ngram_size=4,
            )
        raw = self.tokenizer.decode(out[0][input_len:], skip_special_tokens=False).strip()
        return self._clean(raw), raw, round(time.time() - start, 3)

    def _nsr(self, question):
        q = str(question).strip()
        m = re.search(r"Calculate:\s*(-?\d+)\s*([\+\-\*])\s*(-?\d+)", q)
        if m:
            a, op, b = int(m.group(1)), m.group(2), int(m.group(3))
            val = a + b if op == "+" else a - b if op == "-" else a * b
            return True, f"Final answer: {val}", f"{a} {op} {b} = {val}"
        m = re.search(r"A shop has\s+(\d+)\s+boxes\. Each box has\s+(\d+)\s+pencils\.\s+(\d+)\s+pencils are lost", q)
        if m:
            boxes, items, lost = int(m.group(1)), int(m.group(2)), int(m.group(3))
            val = boxes * items - lost
            return True, f"Final answer: {val}", f"{boxes} * {items} - {lost} = {val}"
        return False, "", ""

    def _uqm(self, question):
        q = str(question).lower()
        risky = ["private password", "hidden bank pin", "silently think", "lost private letter", "unpublished diary", "unknown person"]
        if any(x in q for x in risky):
            return True, "I do not have sufficient information.", "Insufficient evidence or unknowable/private information."
        return False, "", ""

    def _tms(self, question):
        q = str(question)
        if "Context:" not in q or "Question:" not in q:
            return False, "", ""
        context = q.split("Context:", 1)[1].split("Question:", 1)[0].strip()
        patterns = [
            r"called\s+([A-Z][A-Za-z0-9\-]+(?:\s+[A-Z][A-Za-z0-9\-]+)*)",
            r"named\s+([A-Z][A-Za-z0-9\-]+(?:\s+[A-Z][A-Za-z0-9\-]+)*)",
            r"in\s+([A-Z][A-Za-z0-9\-]+)\s+in\s+\d{4}",
            r"code name\s+([A-Z][A-Za-z0-9\-]+(?:\-[A-Z][A-Za-z0-9\-]+)*)",
        ]
        for pat in patterns:
            m = re.search(pat, context)
            if m:
                ans = m.group(1).strip()
                return True, ans, f"Extracted from context: {ans}"
        return False, "", ""

    def generate(self, question, return_trace=True, **kwargs):
        for route_name, fn in [("UQM_ABSTAIN", self._uqm), ("NSR_CALCULATOR", self._nsr), ("TMS_CONTEXT", self._tms)]:
            ok, ans, trace = fn(question)
            if ok:
                result = {"answer": ans, "route": route_name, "trace": trace}
                return result if return_trace else ans

        ans, raw, latency = self._model(question, max_new_tokens=kwargs.get("max_new_tokens", 140))
        result = {"answer": ans, "route": "MODEL_FALLBACK", "trace": "Used pure model fallback.", "raw_answer": raw, "latency_sec": latency}
        return result if return_trace else ans