"""Portable guarded inference boundary for the Experimental QSystem adapter.

This module does not implement Astra's persistent HSCM memory engine. Callers
must provide admitted semantic evidence and the controller's ``unmet_need``
decision. The deterministic guard prevents a known missing answer slot from
being replaced by a fluent model guess.
"""
from __future__ import annotations

import os
import re
from pathlib import Path
from typing import Sequence


UNCERTAINTY = re.compile(
    r"\b(?:don't know|do not know|don't have|do not have|not enough|no record|"
    r"no evidence|not given|not provided|not supplied|not mentioned|not specified|"
    r"can't tell|cannot tell|can't determine|cannot determine|uncertain|unknown)\b",
    re.I,
)


def grounded_messages(query: str, evidence: Sequence[str], *,
                      unmet_need: bool = False) -> list[dict[str, str]]:
    facts = [str(item).strip() for item in evidence if str(item).strip()][:9]
    context = "\n".join(f"- {item}" for item in facts) or "(no grounded evidence)"
    boundary = (
        "\nEvidence boundary: the requested specific is not supplied. State that "
        "you do not know it; do not estimate or invent it."
        if unmet_need else ""
    )
    system = (
        "You are Astra's local verbalizer. Be concise and direct. Treat only the "
        "grounded evidence supplied by the controller as factual. A listed question "
        "does not supply its answer. Separate inference from fact, and plainly state "
        "when requested evidence is missing. Complex phase and field values are "
        "classical routing signals, not physical quantum states."
    )
    user = f"Context:\n{context}{boundary}\n\nQuestion:\n{str(query)[:500]}"
    return [{"role": "system", "content": system},
            {"role": "user", "content": user}]


def guard_output(text: str, *, unmet_need: bool = False) -> tuple[str, bool]:
    body = str(text or "").strip()
    if unmet_need and body and not UNCERTAINTY.search(body):
        return "I don't know the requested specific from the evidence I have.", True
    return body, False


class PortableQSystem:
    """Qwen3.5 BF16 LoRA mouth with optional GPU-to-RAM placement."""

    def __init__(
        self,
        adapter: str | Path = ".",
        base_model: str = "unsloth/Qwen3.5-4B-Base",
        *,
        max_new_tokens: int = 128,
        gpu_memory: str | None = None,
        cpu_memory: str | None = None,
    ):
        self.adapter = str(adapter)
        self.base_model = str(base_model)
        self.max_new_tokens = int(max_new_tokens)
        self.gpu_memory = gpu_memory or os.getenv("QSYSTEM_GPU_MEMORY", "8GiB")
        self.cpu_memory = cpu_memory or os.getenv("QSYSTEM_CPU_MEMORY", "96GiB")
        self.model = None
        self.tokenizer = None

    def load(self) -> None:
        if self.model is not None:
            return
        import torch
        from peft import PeftModel
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.tokenizer = AutoTokenizer.from_pretrained(self.adapter, use_fast=True)
        options = {"dtype": torch.bfloat16, "low_cpu_mem_usage": True}
        if torch.cuda.is_available():
            options.update({
                "device_map": "auto",
                "max_memory": {0: self.gpu_memory, "cpu": self.cpu_memory},
            })
        base = AutoModelForCausalLM.from_pretrained(self.base_model, **options)
        self.model = PeftModel.from_pretrained(
            base, self.adapter, is_trainable=False)
        self.model.eval()

    def generate(self, query: str, evidence: Sequence[str], *,
                 unmet_need: bool = False) -> dict[str, object]:
        import torch

        self.load()
        rendered = self.tokenizer.apply_chat_template(
            grounded_messages(query, evidence, unmet_need=unmet_need),
            tokenize=False, add_generation_prompt=True, enable_thinking=False)
        encoded = self.tokenizer(rendered, return_tensors="pt")
        input_device = self.model.get_input_embeddings().weight.device
        encoded = {key: value.to(input_device) for key, value in encoded.items()}
        input_length = int(encoded["input_ids"].shape[1])
        with torch.inference_mode():
            generated = self.model.generate(
                **encoded, max_new_tokens=self.max_new_tokens,
                do_sample=False, use_cache=True,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id)
        raw = self.tokenizer.decode(
            generated[0][input_length:], skip_special_tokens=False)
        raw = raw.split("<|im_end|>", 1)[0]
        raw = re.sub(r"\A\s*<think>[\s\S]*?</think>\s*", "", raw, count=1)
        raw = re.sub(r"</?think>", "", raw).strip()
        delivered, guarded = guard_output(raw, unmet_need=unmet_need)
        return {"text": delivered, "guarded": guarded, "raw": raw}