"""Qwen-0.5B reasoning router -- Tier 3 for ambiguous requests.

Uses Qwen2.5-0.5B-Instruct via Ollama to reason about:
- What the request actually needs
- Which model best fits the task
- Whether local or cloud is appropriate

Only invoked for the ~5% of requests where kNN (<0.85 confidence)
AND ModernBERT (<0.80 confidence) can't decide.
"""

from __future__ import annotations

import json
import logging
import re
import time
from typing import Any, cast

import httpx
from pydantic import BaseModel

logger = logging.getLogger(__name__)

QWEN_MODEL = "qwen2.5:0.5b"  # Q4_K_M quantization via Ollama
QWEN_AVAILABLE: bool = True  # Module is functional; runtime availability via initialize()

ROUTING_PROMPT = """You are a routing assistant. Given a user request, determine the best model to handle it.

Available models:
{available_models}

User request:
{user_request}

Analyze:
1. Task type (code, creative, research, chat, analysis)
2. Complexity (simple, moderate, complex)
3. Required capabilities (reasoning, coding, knowledge, creativity)
4. Privacy sensitivity (contains PII? medical data? credentials?)

Respond in JSON:
{{
  "task_type": "...",
  "complexity": "...",
  "recommended_tier": "local_small|local_large|cloud",
  "recommended_model": "...",
  "reasoning": "...",
  "privacy_sensitive": false
}}"""


class QwenRoutingResult(BaseModel):
    """Result from the Qwen Tier 3 reasoning router."""

    task_type: str
    complexity: str
    recommended_tier: str  # "local_small", "local_large", "cloud"
    recommended_model: str
    reasoning: str
    privacy_sensitive: bool
    confidence: float  # 0-1
    latency_ms: float
    tier: str = "qwen"  # Always "qwen" for this router


def _uncertain_result(latency_ms: float, reason: str) -> QwenRoutingResult:
    """Return an uncertain result for fallback handling."""
    return QwenRoutingResult(
        task_type="unknown",
        complexity="moderate",
        recommended_tier="cloud",
        recommended_model="",
        reasoning=reason,
        privacy_sensitive=False,
        confidence=0.0,
        latency_ms=latency_ms,
    )


def _extract_json_from_response(text: str) -> dict[str, Any] | None:
    """Extract JSON from model response, handling markdown code blocks."""
    text = text.strip()
    # Try to find ```json ... ``` block
    json_block = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text, re.IGNORECASE)
    if json_block:
        try:
            return cast("dict[str, Any]", json.loads(json_block.group(1).strip()))
        except json.JSONDecodeError:
            pass
    # Try parsing the whole response
    try:
        return cast("dict[str, Any]", json.loads(text))
    except json.JSONDecodeError:
        pass
    # Try to find a JSON object in the text
    obj_match = re.search(r"\{[\s\S]*\}", text)
    if obj_match:
        try:
            return cast("dict[str, Any]", json.loads(obj_match.group(0)))
        except json.JSONDecodeError:
            pass
    return None


def _best_effort_extract(text: str) -> dict[str, Any]:
    """Extract fields from text when JSON parse fails."""
    result: dict[str, Any] = {
        "task_type": "unknown",
        "complexity": "moderate",
        "recommended_tier": "cloud",
        "recommended_model": "",
        "reasoning": f"Parse fallback: {text[:200]}",
        "privacy_sensitive": False,
    }
    # Heuristic: look for quoted strings after key names
    for key in ["task_type", "complexity", "recommended_tier", "recommended_model"]:
        pat = re.compile(rf'"{key}"\s*:\s*"([^"]*)"', re.IGNORECASE)
        m = pat.search(text)
        if m:
            result[key] = m.group(1)
    # Look for "privacy_sensitive": true
    if re.search(r'"privacy_sensitive"\s*:\s*true', text, re.IGNORECASE):
        result["privacy_sensitive"] = True
    return result


class QwenRouter:
    """Tier 3 reasoning router using Qwen2.5-0.5B via Ollama.

    Decision flow:
    1. Format routing prompt with available models
    2. Call Ollama /api/generate with Qwen
    3. Parse JSON response
    4. Map to routing decision

    Fallbacks:
    - If Qwen not installed: return UNCERTAIN (let rule-based handle)
    - If Ollama unreachable: return UNCERTAIN
    - If response timeout (>500ms): return UNCERTAIN
    - If JSON parse fails: extract best-effort from text
    """

    ROUTE_TIMEOUT: float = 0.5  # 500ms for routing decision (PRD)
    INIT_TIMEOUT: float = 5.0  # 5s for warm-up / model check

    def __init__(self, ollama_host: str = "http://localhost:11434") -> None:
        self._host = ollama_host.rstrip("/")
        self._model = QWEN_MODEL
        self._available = False

    async def initialize(self) -> bool:
        """Check if Qwen model is available in Ollama."""
        try:
            async with httpx.AsyncClient(timeout=self.INIT_TIMEOUT) as client:
                r = await client.get(f"{self._host}/api/tags")
                r.raise_for_status()
                data = r.json()
                models = data.get("models", [])
                names = [m.get("name", "").split(":")[0] for m in models if m.get("name")]
                # Check for qwen2.5 (tag may be qwen2.5:0.5b or qwen2.5:0.5b-q4_k_m)
                self._available = any("qwen2.5" in n or "qwen" in n for n in names)
                if not self._available:
                    logger.info("Qwen model not found in Ollama. Run: ollama pull %s", self._model)
                return self._available
        except Exception as e:
            logger.warning("QwenRouter: Ollama unreachable during initialize: %s", e)
            self._available = False
            return False

    async def route(self, prompt: str, available_models: list[str]) -> QwenRoutingResult:
        """Reason about a request and return routing recommendation."""
        start = time.perf_counter()

        if not available_models:
            elapsed = (time.perf_counter() - start) * 1000
            return _uncertain_result(
                elapsed,
                "No available models provided; cannot recommend. Escalating to rule-based.",
            )

        models_str = ", ".join(available_models)
        routing_prompt = ROUTING_PROMPT.format(
            available_models=models_str,
            user_request=prompt[:2000],  # Cap to avoid token overflow
        )

        try:
            async with httpx.AsyncClient(timeout=self.ROUTE_TIMEOUT) as client:
                resp = await client.post(
                    f"{self._host}/api/generate",
                    json={
                        "model": self._model,
                        "prompt": routing_prompt,
                        "stream": False,
                        "options": {"num_predict": 256},
                    },
                )
                resp.raise_for_status()
                data = resp.json()
        except httpx.TimeoutException:
            elapsed = (time.perf_counter() - start) * 1000
            logger.warning("QwenRouter: route timeout (>%sms)", int(self.ROUTE_TIMEOUT * 1000))
            return _uncertain_result(elapsed, "Ollama timeout; escalating to rule-based.")
        except httpx.ConnectError:
            elapsed = (time.perf_counter() - start) * 1000
            logger.warning("QwenRouter: Ollama unreachable")
            return _uncertain_result(elapsed, "Ollama unreachable; escalating to rule-based.")
        except Exception as e:
            elapsed = (time.perf_counter() - start) * 1000
            logger.warning("QwenRouter: route error: %s", e, exc_info=True)
            return _uncertain_result(elapsed, f"Qwen routing failed: {e}. Escalating.")

        elapsed = (time.perf_counter() - start) * 1000
        response_text = data.get("response", "")
        if not response_text:
            return _uncertain_result(elapsed, "Empty Ollama response; escalating.")

        parsed = _extract_json_from_response(response_text)
        used_best_effort = False
        if parsed is None:
            parsed = _best_effort_extract(response_text)
            used_best_effort = True
            logger.debug("QwenRouter: JSON parse failed, used best-effort extraction")

        # Map parsed fields to QwenRoutingResult
        task_type = str(parsed.get("task_type", "unknown"))
        complexity = str(parsed.get("complexity", "moderate"))
        recommended_tier = str(parsed.get("recommended_tier", "cloud"))
        recommended_model = str(parsed.get("recommended_model", ""))
        reasoning = str(parsed.get("reasoning", ""))
        privacy_sensitive = bool(parsed.get("privacy_sensitive", False))

        # Validate recommended_tier
        if recommended_tier not in ("local_small", "local_large", "cloud"):
            recommended_tier = "cloud"

        # Pick a model if recommended isn't in available_models
        if recommended_model and recommended_model not in available_models:
            # Try to match by base name (e.g. qwen2.5:7b matches qwen2.5:7b-q4)
            rec_base = recommended_model.lower().split(":")[0]
            for m in available_models:
                if m.lower().split(":")[0] == rec_base:
                    recommended_model = m
                    break
            else:
                recommended_model = available_models[0] if available_models else ""

        if not recommended_model and available_models:
            recommended_model = available_models[0]

        # Confidence: high if we got valid JSON; lower for best-effort parse
        confidence = 0.65 if used_best_effort else 0.85

        return QwenRoutingResult(
            task_type=task_type,
            complexity=complexity,
            recommended_tier=recommended_tier,
            recommended_model=recommended_model,
            reasoning=reasoning,
            privacy_sensitive=privacy_sensitive,
            confidence=confidence,
            latency_ms=elapsed,
        )

    async def ensure_model(self) -> bool:
        """Check if model exists, suggest pulling if not."""
        if self._available:
            return True
        available = await self.initialize()
        if not available:
            logger.info("Pull Qwen model with: ollama pull %s", self._model)
        return available
