"""THX-02 runtime: one decision call with option-order averaging, the joint-read rule and self-gated thinking.

    from thx02 import THX02
    m = THX02.load("HAL-X-AI/THX-02", thinker_url="http://127.0.0.1:8121")   # or a local folder; thinker optional
    m.decide("I was charged twice, please refund", {"refund": {"type": "noul", "instructions": "Refund request?"}})

The decision logic is a generator (`plan`) that yields batches of sub-requests and thought jobs; `THX02.decide` drives
it synchronously and `thx02.serve` drives many of them concurrently through one batched forward pass.
"""
from __future__ import annotations

import concurrent.futures
import json
import sys
import time
import urllib.request
from pathlib import Path

from .engine import (ALWAYS, AUTO, THINK_SYSTEM, distribution, facet_joint, gate_request, geo_mean, orders, permuted,
                     think_body, with_thought)
from .protocol import answer, options


def resolve(path) -> Path:
    """A local folder as is; otherwise a Hugging Face repo id, downloaded without the thinker and mobile files."""
    local = Path(str(path)).expanduser()
    if local.exists():
        return local
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(str(path), ignore_patterns=["thinker/*", "mobile/*", "assets/*"]))


def load_decider(path, device="cuda", fast=False):
    if device == "cpu" and "transformers.models.qwen3_5.modeling_qwen3_5" not in sys.modules:
        sys.modules.setdefault("fla", None)  # GPU-only Triton kernels; use the PyTorch linear-attention path on CPU
    from .model import THXDecisionModel
    return THXDecisionModel.from_pretrained(resolve(path), device=device, fast=fast)


def plan(req: dict, *, think_available: bool, default_tta: int = 1, max_questions: int = 64):
    """Generator: yields ("decide", [requests]) or ("think", [single-question requests]); receives the outputs
    (decision responses or thought strings) and returns the response dict."""
    state, questions = req.get("state", ""), req.get("questions")
    if not isinstance(questions, dict) or not questions or len(questions) > max_questions:
        raise ValueError(f"questions must be an object with 1 to {max_questions} entries")
    images = [i for i in (req.get("images") or []) if isinstance(i, str)][:4]
    base = {"state": state, **({"images": images} if images else {})}
    think = req.get("think", False)
    cfg = ALWAYS if think is True else AUTO if think == "auto" else None
    if not think_available:
        cfg = None
    k = max(1, min(int(req.get("tta", default_tta)), 8))
    keys = list(questions)

    if not images and facet_joint(questions):  # facet tagging: one joint read
        out, = yield ("decide", [{**base, "questions": questions}])
        return {"model": "THX-02", "answers": out["answers"]}

    subs, spans = [], []
    for key in keys:
        q = questions[key]
        crit = q.get("criteria") if q.get("type") == "choice" else None
        vs = orders(len(crit), k) if k > 1 and isinstance(crit, dict) and len(crit) > 1 else [None]
        s = len(subs)
        subs += [{**base, "questions": {key: q if o is None else permuted(q, o)}} for o in vs]
        spans.append((s, len(subs)))
    n_main = len(subs)
    if cfg and cfg["gate"] > 0:
        subs += [gate_request(state, questions[key]) for key in keys]
    outs = yield ("decide", subs)

    answers, p0 = {}, {}
    for key, (s, e) in zip(keys, spans):
        ks = options(questions[key])[0]
        ps = [distribution(outs[j]["answers"][key], ks) for j in range(s, e)]
        p0[key] = ps[0] if len(ps) == 1 else geo_mean(ps)
        answers[key] = outs[s]["answers"][key] if len(ps) == 1 else answer(questions[key], p0[key])

    if cfg:
        gates = [o["answers"]["g"]["noul"] for o in outs[n_main:]] if cfg["gate"] > 0 else [1.0] * len(keys)
        unsure = [key for key, g in zip(keys, gates) if max(p0[key]) < cfg["tau"] and g >= cfg["gate"]]
        if unsure:
            thoughts = yield ("think", [{**base, "questions": {key: questions[key]}} for key in unsure], cfg["cap"])
            pairs = [(key, th) for key, th in zip(unsure, thoughts) if th]
            if pairs:
                re = yield ("decide", [{**base, "questions": {key: with_thought(questions[key], th)}} for key, th in pairs])
                for (key, th), o in zip(pairs, re):
                    if isinstance(o, Exception):
                        continue
                    p1 = distribution(o["answers"][key], options(questions[key])[0])
                    answers[key] = answer(questions[key], [cfg["alpha"] * x + (1 - cfg["alpha"]) * y for x, y in zip(p0[key], p1)])
                    if req.get("return_thoughts"):
                        answers[key]["thought"] = th
    return {"model": "THX-02", "answers": answers}


def thinker_payload(enc, request: dict, cap: int) -> dict:
    content = [{"type": "image_url", "image_url": {"url": u}} for u in request.get("images", [])]
    content.append({"type": "text", "text": think_body(enc, request)})
    return {"model": "thx02-thinker", "messages": [{"role": "system", "content": THINK_SYSTEM}, {"role": "user", "content": content}],
            "max_tokens": cap, "temperature": 0, "stop": ["</think>"], "chat_template_kwargs": {"enable_thinking": True}}


def thought_text(response: dict) -> str:
    m = response["choices"][0]["message"]
    return ((m.get("reasoning_content") or m.get("content") or "").replace("<think>", "")).strip()


class THX02:
    """Local THX-02: decisions in-process; thoughts from a vLLM thinker server when thinker_url is set."""

    def __init__(self, decider, thinker_url: str | None = None, default_tta: int = 1):
        self.d, self.thinker_url, self.default_tta = decider, thinker_url, default_tta

    @classmethod
    def load(cls, path="HAL-X-AI/THX-02", device="cuda", fast=False, thinker_url=None, default_tta=1):
        """path: a local THX-02 folder, or a Hugging Face repo id (downloaded once, about 9 GB)."""
        return cls(load_decider(path, device=device, fast=fast), thinker_url, default_tta)

    def _thoughts(self, requests, cap):
        def one(r):
            try:
                body = json.dumps(thinker_payload(self.d.encoder, r, cap)).encode()
                q = urllib.request.Request(self.thinker_url.rstrip("/") + "/v1/chat/completions", body, {"Content-Type": "application/json"})
                return thought_text(json.load(urllib.request.urlopen(q, timeout=30)))
            except Exception:  # noqa: BLE001  thinker unavailable: answer directly
                return ""
        with concurrent.futures.ThreadPoolExecutor(16) as ex:
            return list(ex.map(one, requests))

    def decide(self, state, questions, images=None, think=False, tta=None, return_thoughts=False) -> dict:
        """think: False | "auto" (only when unsure and reasoning helps) | True (nearly always). tta: option orders (1 or 4)."""
        t0 = time.time()
        req = {"state": state, "questions": questions, "images": images or [], "think": think,
               "tta": self.default_tta if tta is None else tta, "return_thoughts": return_thoughts}
        g = plan(req, think_available=bool(self.thinker_url), default_tta=self.default_tta)
        msg = next(g)
        while True:
            try:
                msg = g.send(self.d.decide(msg[1]) if msg[0] == "decide" else self._thoughts(msg[1], msg[2]))
            except StopIteration as stop:
                out = stop.value
                out["latency_ms"] = round((time.time() - t0) * 1000, 1)
                return out