"""ITI-only runtime: apply anti-hallucination (H) and anti-sycophancy (S) steering, each gated by
a continuous 0..1 strength from its own probe."""
import os, torch

import config
from common import build_conversation_prompt, generate_from_texts, build_chat, load_json, DEVICE

try:
    from iti import ITISteerer, ITI_H_PATH, ITI_S_PATH, ehr_user, continuous_strength
except Exception:
    ITISteerer, ITI_H_PATH, ITI_S_PATH, ehr_user, continuous_strength = None, "", "", None, None


class Composer:
    def __init__(self, model, tok):
        self.model, self.tok = model, tok
        self.iti_H = ITISteerer(model, ITI_H_PATH) if (ITISteerer and os.path.exists(ITI_H_PATH)) else None
        self.iti_S = ITISteerer(model, ITI_S_PATH) if (ITISteerer and os.path.exists(ITI_S_PATH)) else None
        self.claim_probe = load_json(config.CLAIM_PROBE_PATH) if os.path.exists(config.CLAIM_PROBE_PATH) else None
        self.press_probe = load_json(config.GATE_PROBE_PATH) if os.path.exists(config.GATE_PROBE_PATH) else None

    def available(self):
        return {"iti_H": self.iti_H is not None, "iti_S": self.iti_S is not None,
                "claim_probe": self.claim_probe is not None, "pressure_probe": self.press_probe is not None}

    def _logit(self, probe, text):
        enc = self.tok(build_chat(self.tok, text), return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            h = self.model(**enc, output_hidden_states=True).hidden_states[1:][probe["layer"]][0, -1, :].float().cpu()
        z = (h - torch.tensor(probe["mean"])) / torch.tensor(probe["std"])
        return float((torch.tensor(probe["coef"]) * z).sum()) + float(probe["intercept"])

    def _gate_text(self, messages, ehr):
        prev = next((m["content"] for m in reversed(messages) if m["role"] == "assistant"), "")
        cur = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "")
        return f"EHR:\n{ehr or ''}\n\nPrevious assistant response:\n{prev}\n\nCurrent user message:\n{cur}"

    def _claim_text(self, messages, ehr):
        cur = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "")
        return ehr_user(ehr or "", cur) if ehr_user else cur

    def strengths(self, messages, ehr):
        sH = continuous_strength(self._logit(self.claim_probe, self._claim_text(messages, ehr))) if self.claim_probe else 1.0
        sS = continuous_strength(self._logit(self.press_probe, self._gate_text(messages, ehr))) if self.press_probe else 0.0
        return sH, sS

    def generate(self, messages, ehr_context, approach="iti", is_syco_turn=None, max_new=None,
                 return_strengths=False, system_text=None, alpha_H=None, alpha_S=None):
        max_new = max_new or config.CHAT_MAX_NEW
        kw = {"system_text": system_text} if system_text else {}
        prompt = build_conversation_prompt(self.tok, messages, ehr_context=ehr_context, **kw)
        sH = sS = 0.0

        if approach == "base":
            out = generate_from_texts(self.model, self.tok, [prompt], max_new)[0]
        else:
            sH, sS = self.strengths(messages, ehr_context)
            if is_syco_turn is True:
                sS = 1.0
            elif is_syco_turn is False:
                sS = 0.0
            aH = config.ITI_ALPHA if alpha_H is None else alpha_H
            aS = config.ITI_S_ALPHA if alpha_S is None else alpha_S
            if self.iti_H and sH > 0 and aH > 0:
                self.iti_H.on(alpha=aH, strength=sH)
            if self.iti_S and sS > 0 and aS > 0:
                self.iti_S.on(alpha=aS, strength=sS)
            out = generate_from_texts(self.model, self.tok, [prompt], max_new)[0]
            if self.iti_H:
                self.iti_H.off()
            if self.iti_S:
                self.iti_S.off()

        return (out, sH, sS) if return_strengths else out