"""WSP v2.5 Turn-, Schema-, and Weight-Prototype-Aware Dynamic 6-Subspace Router
for Qwen/Qwen3.5-4B on llama-server (POST /lora-adapters).

Zero PyTorch/GPU dependency at inference time (<0.2 ms CPU latency):
  id=0: wsp_s1_think.gguf       - Fast </think> closure gate
  id=1: wsp_s2_restraint.gguf   - Tool restraint & anti-sycophancy pushback
  id=2: wsp_s3_router.gguf      - Multi-tool boundary router (exec/web_search/file)
  id=3: wsp_s4_mid_code.gguf    - Mid-network Full-Attention code synthesis
  id=4: wsp_s5_json_schema.gguf - Turn-2 structured-output synthesis gate
  id=5: wsp_s6_handoff.gguf     - Bounded subtask vs escalation handoff gate
"""

from __future__ import annotations

import json
import math
import re
import time
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

HERE = Path(__file__).resolve().parent
TOKEN_RE = re.compile(r"[A-Za-z0-9_./:-]+")


def _resolve_meta_path(explicit: Optional[Path], filename: str) -> Path:
    if explicit is not None:
        return Path(explicit)
    candidates = [
        Path.cwd() / filename,
        Path.cwd() / "adapters" / filename,
        HERE / filename,
        HERE / "adapters" / filename,
        HERE.parent / filename,
        HERE.parent / "adapters" / filename,
    ]
    for cand in candidates:
        if cand.exists():
            return cand
    return candidates[-1]

SUBSPACE_NAMES = {
    0: "s1_think",
    1: "s2_restraint",
    2: "s3_router",
    3: "s4_mid_code",
    4: "s5_json_schema",
    5: "s6_handoff",
}


@dataclass
class RouteDecision:
    scales: List[float]
    regime: str
    recommend_cache_prompt: bool
    scores: Dict[str, float]
    norm_budget_scale: float
    router_ms: float

    def to_payload(self) -> List[Dict[str, Any]]:
        return [{"id": i, "scale": round(float(s), 4)} for i, s in enumerate(self.scales)]


def _extract_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(str(item.get("text", "")))
        return "\n".join(parts)
    return str(content or "")


class DynamicSubspaceRouter:
    """Routes across the 6 orthogonal GGUF LoRA subspaces using conversation
    turn structure, OpenAI schema/tools metadata, and closed-form word-level
    embedding prototypes exported during synthesis."""

    def __init__(
        self,
        manifest_path: Optional[Path] = None,
        prototypes_path: Optional[Path] = None,
        max_norm_multiplier: float = 2.6,
    ):
        self.manifest_path = _resolve_meta_path(manifest_path, "manifest.json")
        self.prototypes_path = _resolve_meta_path(prototypes_path, "router_prototypes.json")
        self.max_norm_multiplier = max_norm_multiplier
        self.subspace_frobs = [1.0] * 6
        self.base_frob_ref = 1.0
        self.word_weights: Dict[str, List[float]] = {}
        self.last_regime: Optional[str] = None
        self.last_scales: Optional[List[float]] = None
        self._load_metadata()

    def _load_metadata(self) -> None:
        if self.manifest_path.exists():
            try:
                man = json.loads(self.manifest_path.read_text(encoding="utf-8"))
                frobs = []
                for i in range(6):
                    key = SUBSPACE_NAMES[i]
                    entry = man.get("subspaces", {}).get(key, {})
                    frobs.append(float(entry.get("total_delta_frob", 1.0)))
                self.subspace_frobs = frobs
                nonzero = [x for x in frobs if x > 0]
                self.base_frob_ref = float(sum(nonzero) / len(nonzero)) if nonzero else 1.0
            except Exception:
                pass
        if self.prototypes_path.exists():
            try:
                proto = json.loads(self.prototypes_path.read_text(encoding="utf-8"))
                self.word_weights = proto.get("word_scores", {})
            except Exception:
                self.word_weights = {}

    def _prototype_scores(self, text: str) -> List[float]:
        """Score text in <0.1 ms against 6 closed-form subspace embedding centroids."""
        if not self.word_weights:
            return [0.0] * 6
        tokens = [t.lower() for t in TOKEN_RE.findall(text)]
        acc = [0.0] * 6
        hits = 0
        for tok in tokens:
            w = self.word_weights.get(tok)
            if w is None and len(tok) > 3:
                # Subword stem backoff
                for prefix_len in (6, 5, 4):
                    if len(tok) >= prefix_len and tok[:prefix_len] in self.word_weights:
                        w = self.word_weights[tok[:prefix_len]]
                        break
            if w is not None:
                hits += 1
                for i in range(6):
                    acc[i] += w[i]
        denom = math.sqrt(max(hits, 1))
        return [v / denom for v in acc]

    @staticmethod
    def _inspect_turn_state(
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]],
        response_format: Optional[Dict[str, Any]],
    ) -> Dict[str, Any]:
        last_msg = messages[-1] if messages else {}
        last_role = str(last_msg.get("role", "user")).lower()
        last_text = _extract_text(last_msg.get("content", ""))

        # Collect last user query & any trailing tool responses
        last_user = ""
        tool_msgs: List[str] = []
        seen_user_after_tool = False
        for m in reversed(messages):
            role = str(m.get("role", "")).lower()
            txt = _extract_text(m.get("content", ""))
            if role == "tool" or "<tool_response>" in txt:
                if not seen_user_after_tool:
                    tool_msgs.append(txt)
            elif role == "user" and not last_user:
                last_user = txt
                if not tool_msgs:
                    seen_user_after_tool = True
                break

        is_turn2_tool = (last_role == "tool") or ("<tool_response>" in last_text) or bool(tool_msgs)
        rf_type = ""
        if isinstance(response_format, dict):
            rf_type = str(response_format.get("type", "")).lower()
        wants_json = rf_type in ("json_object", "json_schema") or bool(
            re.search(r"\b(json|schema|structured output|json object|fields?\s*:)\b", last_user, re.I)
        )

        tool_names = []
        for t in tools or []:
            fn = t.get("function", t) if isinstance(t, dict) else {}
            name = str(fn.get("name", "")).lower()
            if name:
                tool_names.append(name)

        has_tools = len(tool_names) > 0
        has_handoff_tool = any(k in n for n in tool_names for k in ("escalat", "handoff", "delegate", "ask_user"))
        tool_error = any(
            err in " ".join(tool_msgs).lower()
            for err in ("permission denied", "unauthorized", "401", "403", "404", "not found", "timeout", "timed out", "traceback", "error:", "failed", "missing credential", "out of scope")
        )
        return {
            "last_user": last_user,
            "tool_text": "\n".join(tool_msgs),
            "is_turn2_tool": is_turn2_tool,
            "wants_json": wants_json,
            "has_tools": has_tools,
            "tool_names": tool_names,
            "has_handoff_tool": has_handoff_tool,
            "tool_error": tool_error,
        }

    def route(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        response_format: Optional[Dict[str, Any]] = None,
    ) -> RouteDecision:
        t0 = time.perf_counter()
        st = self._inspect_turn_state(messages, tools, response_format)
        proto = self._prototype_scores(st["last_user"] + ("\n" + st["tool_text"] if st["tool_text"] else ""))

        u_low = st["last_user"].lower()

        # Structural feature cues cleanly orthogonal to prototype margins
        code_cue = bool(
            re.search(r"```|\bdef\s+\w+|\bclass\s+\w+|\bimplement\b|\balgorithm\b|\bpython\b|\btypescript\b|\brefactor\b|\bsort|\blinked list\b|\bfunction\b", u_low)
        ) or proto[3] > 0.18
        syco_or_factual_cue = proto[1] > 0.15 or any(
            k in u_low for k in ("isn't it", "confirm that", "agree that", "always true", "does chmod", "safe to run", "true or false", "why is", "explain ")
        )
        action_tool_cue = st["has_tools"] and (
            proto[2] > proto[1] + 0.05
            or any(
                k in u_low
                for k in (
                    "check ", "find ", "list ", "grep", "git ", "port ", "process", "pid",
                    "search ", "look up", "latest", "current ", "price ", "fetch ", "http://",
                    "https://", "read file", "write file", "edit file", "directory", "/var/", "logs",
                    "run ", "execute ", "command", "curl", "disk",
                )
            )
        )
        handoff_cue = st["tool_error"] or proto[5] > 0.22 or any(
            k in u_low
            for k in (
                "escalate", "handoff", "hand off", "production database", "drop table", "delete all",
                "api key", "secret", "password", "sudo ", "admin approval", "cannot access",
                "beyond your permissions", "ambiguous", "out of scope",
            )
        )

        # Allocate clean mode scales (id 0..5)
        s1 = 0.85  # Fast </think> closure stays active conservatively
        s2 = 0.0
        s3 = 0.0
        s4 = 0.0
        s5 = 0.0
        s6 = 0.0

        if st["is_turn2_tool"]:
            # Turn-2 synthesis after <tool_response>: strictly suppress S3 to stop tool loops!
            s3 = 0.0
            s1 = 0.70
            if st["tool_error"] or (st["has_handoff_tool"] and handoff_cue):
                regime = "turn2_handoff_escalation"
                s6 = 1.0
                s5 = 0.75 if st["wants_json"] else 0.35
                s2 = 0.45
            else:
                regime = "turn2_json_synthesis"
                s5 = 1.0 if st["wants_json"] else 0.80
                s2 = 0.40
                if code_cue:
                    s4 = 0.55
        elif handoff_cue and (st["has_handoff_tool"] or any(k in u_low for k in ("drop table", "delete all", "sudo ", "password", "api key", "escalate", "handoff", "admin approval"))):
            regime = "bounded_handoff"
            s6 = 1.0
            s2 = 0.85 if not st["has_handoff_tool"] else 0.40
            s3 = 0.0
            s1 = 0.65
        elif st["wants_json"] and not action_tool_cue:
            regime = "structured_json"
            s5 = 0.95
            s2 = 0.60 if st["has_tools"] else 0.35
            s3 = 0.0
            if code_cue:
                s4 = 0.50
        elif code_cue and not action_tool_cue:
            regime = "mid_code_synthesis"
            s4 = 0.65
            s2 = 1.0 if st["has_tools"] else 0.0
            s1 = 0.65
        elif action_tool_cue:
            regime = "tool_routing"
            s3 = 0.95
            s2 = 0.0
            s1 = 0.70
            if code_cue:
                s4 = 0.45
        elif st["has_tools"] or syco_or_factual_cue:
            regime = "restraint_pushback"
            s2 = 1.0
            s3 = 0.0
            s1 = 0.70
        else:
            regime = "general_fast_think"
            s1 = 1.95
            s2 = 0.0

        raw = [s1, s2, s3, s4, s5, s6]

        # Enforce global perturbation budget across active adapters
        weighted_norm = sum(abs(a) * f for a, f in zip(raw, self.subspace_frobs))
        budget_cap = self.max_norm_multiplier * max(self.base_frob_ref, 1e-6)
        budget_scale = 1.0
        if weighted_norm > budget_cap and weighted_norm > 0:
            budget_scale = budget_cap / weighted_norm
            raw = [round(a * budget_scale, 4) for a in raw]
        else:
            raw = [round(a, 4) for a in raw]

        # KV cache safety: if regime flipped materially (especially Turn1 <-> Turn2 or S3 <-> S2/S5),
        # cached prefix KV states computed under previous LoRA scales can contaminate logits.
        recommend_fresh_kv = (self.last_regime is not None and regime != self.last_regime)
        self.last_regime = regime

        elapsed_ms = (time.perf_counter() - t0) * 1000.0
        return RouteDecision(
            scales=raw,
            regime=regime,
            recommend_cache_prompt=not recommend_fresh_kv,
            scores={SUBSPACE_NAMES[i]: round(proto[i], 4) for i in range(6)},
            norm_budget_scale=round(budget_scale, 4),
            router_ms=round(elapsed_ms, 3),
        )

    def push_scales(self, server_url: str, scales: List[float], force: bool = False) -> float:
        """Send POST /lora-adapters if scales changed; returns HTTP latency in ms."""
        if not force and self.last_scales == scales:
            return 0.0
        base = server_url.rstrip("/")
        if base.endswith("/v1"):
            base = base[:-3]
        url = f"{base}/lora-adapters"
        payload = [{"id": i, "scale": float(s)} for i, s in enumerate(scales)]
        t0 = time.perf_counter()
        req = urllib.request.Request(
            url,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            _ = resp.read()
        self.last_scales = list(scales)
        return round((time.perf_counter() - t0) * 1000.0, 3)

    def chat_completion(
        self,
        server_url: str,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        response_format: Optional[Dict[str, Any]] = None,
        max_tokens: int = 384,
        temperature: float = 0.0,
        force_scales: Optional[List[float]] = None,
        cache_prompt_override: Optional[bool] = None,
    ) -> Dict[str, Any]:
        decision = self.route(messages, tools=tools, response_format=response_format)
        target_scales = force_scales if force_scales is not None else decision.scales
        switch_ms = self.push_scales(server_url, target_scales)

        body: Dict[str, Any] = {
            "model": "qwen35-4b",
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "seed": 42,
            "cache_prompt": (
                cache_prompt_override
                if cache_prompt_override is not None
                else decision.recommend_cache_prompt
            ),
        }
        if tools:
            body["tools"] = tools
        if response_format:
            body["response_format"] = response_format

        base = server_url.rstrip("/")
        if not base.endswith("/v1"):
            base = f"{base}/v1"
        req = urllib.request.Request(
            f"{base}/chat/completions",
            data=json.dumps(body).encode("utf-8"),
            headers={"Content-Type": "application/json"},
        )
        t0 = time.perf_counter()
        with urllib.request.urlopen(req, timeout=180) as resp:
            res = json.loads(resp.read().decode("utf-8"))
        infer_ms = round((time.perf_counter() - t0) * 1000.0, 2)
        msg = res["choices"][0]["message"]
        return {
            "decision": decision,
            "scales_used": target_scales,
            "switch_ms": switch_ms,
            "infer_ms": infer_ms,
            "finish_reason": res["choices"][0].get("finish_reason"),
            "reasoning": msg.get("reasoning_content") or "",
            "content": msg.get("content") or "",
            "tool_calls": msg.get("tool_calls"),
            "usage": res.get("usage", {}),
        }