"""
Thin policy-based query routing client.

**Tier: Advanced Contract (MAY)**

Routes incoming queries to the most appropriate model tier by fetching
a routing policy from the Octomil server and applying simple threshold
comparisons locally.  Complex scoring logic lives server-side.

**Online mode**: POST /api/v1/route/query for full server-side routing.
**Offline mode**: Apply cached policy with simple word-count / keyword
thresholds.  Falls back to an embedded default policy on first use.

Backward-compatible public API:
- ``ModelInfo``, ``RoutingDecision``, ``DecomposedRoutingDecision``,
  ``DeterministicResult`` dataclasses
- ``QueryRouter`` class with ``route()``, ``route_decomposed()``,
  ``get_fallback()``
- ``check_deterministic()`` and ``assign_tiers()`` module-level functions
"""

from __future__ import annotations

import json
import logging
import os
import re
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional, Union

from .errors import OctomilError, OctomilErrorCode

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Model metadata
# ---------------------------------------------------------------------------

TIER_ORDER: list[str] = ["fast", "balanced", "quality"]


@dataclass
class ModelInfo:
    """Metadata about a loaded model used for routing decisions."""

    name: str
    tier: str = "balanced"
    param_b: float = 0.0
    loaded: bool = True

    @property
    def tier_index(self) -> int:
        """Numeric rank of the tier (higher = more capable)."""
        try:
            return TIER_ORDER.index(self.tier)
        except ValueError:
            return 1


# ---------------------------------------------------------------------------
# Deterministic result
# ---------------------------------------------------------------------------


@dataclass
class DeterministicResult:
    """Result from a deterministic (no-model) computation."""

    answer: str
    method: str
    confidence: float = 1.0


# ---------------------------------------------------------------------------
# Routing decision
# ---------------------------------------------------------------------------


@dataclass
class RoutingDecision:
    """Result of a routing decision with metadata for telemetry."""

    model_name: str
    complexity_score: float
    tier: str
    strategy: str
    fallback_chain: list[str] = field(default_factory=list)
    deterministic_result: Optional[DeterministicResult] = None


@dataclass
class DecomposedRoutingDecision:
    """Routing decision for a decomposed multi-task query."""

    sub_decisions: list[RoutingDecision]
    tasks: list  # list[SubTask] — typed loosely to avoid circular import
    original_query: str
    decomposed: bool = True


# ---------------------------------------------------------------------------
# Routing policy (fetched from server, cached locally)
# ---------------------------------------------------------------------------

_MATH_PATTERN = re.compile(
    r"^[\s\d+\-*/^().%]+$"
    r"|[\d]+\s*[+\-*/^%]\s*[\d]"
    r"|\b(sqrt|abs|log|sin|cos|tan|ceil|floor)\s*\("
    r"|[\d.]+\s*%\s*of\s+[\d.]+",
)

_DEFAULT_POLICY: dict[str, Any] = {
    "version": 1,
    "thresholds": {"fast_max_words": 0, "quality_min_words": 999999},
    "complex_indicators": [],
    "deterministic_enabled": True,
    "ttl_seconds": 0,
}


@dataclass
class RoutingPolicy:
    """Cached routing policy from the server."""

    version: int
    fast_max_words: int
    quality_min_words: int
    complex_indicators: list[str]
    deterministic_enabled: bool
    ttl_seconds: int
    fetched_at: float = 0.0
    etag: str = ""
    # --- iOS-aligned scoring weights (new) ---
    length_weight: float = 0.5
    indicator_weight: float = 0.5
    fast_threshold: float = 0.3
    quality_threshold: float = 0.7
    indicator_normalizor: float = 3.0

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> RoutingPolicy:
        thresholds = d.get("thresholds", {})

        return cls(
            version=d.get("version", 1),
            fast_max_words=thresholds.get("fast_max_words", 10),
            quality_min_words=thresholds.get("quality_min_words", 50),
            complex_indicators=d.get("complex_indicators", []),
            deterministic_enabled=d.get("deterministic_enabled", True),
            ttl_seconds=d.get("ttl_seconds", 3600),
            fetched_at=d.get("fetched_at", 0.0),
            etag=d.get("etag", ""),
            length_weight=float(d.get("length_weight", 0.5)),
            indicator_weight=float(d.get("indicator_weight", 0.5)),
            fast_threshold=float(d.get("fast_threshold", 0.3)),
            quality_threshold=float(d.get("quality_threshold", 0.7)),
            indicator_normalizor=float(d.get("indicator_normalizor", 3.0)),
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "version": self.version,
            "thresholds": {
                "fast_max_words": self.fast_max_words,
                "quality_min_words": self.quality_min_words,
            },
            "complex_indicators": self.complex_indicators,
            "deterministic_enabled": self.deterministic_enabled,
            "ttl_seconds": self.ttl_seconds,
            "fetched_at": self.fetched_at,
            "etag": self.etag,
            "length_weight": self.length_weight,
            "indicator_weight": self.indicator_weight,
            "fast_threshold": self.fast_threshold,
            "quality_threshold": self.quality_threshold,
            "indicator_normalizor": self.indicator_normalizor,
        }

    @property
    def is_expired(self) -> bool:
        if self.fetched_at == 0.0:
            return True
        return (time.time() - self.fetched_at) > self.ttl_seconds


# ---------------------------------------------------------------------------
# Policy client — fetches and caches policy from server
# ---------------------------------------------------------------------------

_CACHE_DIR = Path(os.environ.get("OCTOMIL_CACHE_DIR", Path.home() / ".octomil"))
_POLICY_FILE = "routing_policy.json"


class PolicyClient:
    """Fetches routing policy from the server and caches to disk."""

    def __init__(
        self,
        api_base: str = "https://api.octomil.com/api/v1",
        api_key: str = "",
    ) -> None:
        self.api_base = api_base.rstrip("/")
        self.api_key = api_key
        self._policy: Optional[RoutingPolicy] = None

    def get_policy(self) -> RoutingPolicy:
        """Return the current policy, refreshing from server if expired."""
        if self._policy is not None and not self._policy.is_expired:
            return self._policy

        # Try loading from disk cache
        if self._policy is None:
            self._policy = self._load_from_disk()

        # If still valid after loading from disk, return it
        if self._policy is not None and not self._policy.is_expired:
            return self._policy

        # Try fetching from server
        fetched = self._fetch_from_server()
        if fetched is not None:
            self._policy = fetched
            self._save_to_disk(fetched)
            return self._policy

        # Use disk cache even if expired, or fall back to defaults
        if self._policy is not None:
            logger.debug("Using expired cached policy (server unreachable)")
            return self._policy

        logger.debug("Using default embedded policy (no cache, server unreachable)")
        self._policy = RoutingPolicy.from_dict(_DEFAULT_POLICY)
        return self._policy

    def _fetch_from_server(self) -> Optional[RoutingPolicy]:
        """Fetch policy from GET /api/v1/route/policy."""
        try:
            import httpx
        except ImportError:
            logger.debug("httpx not available — cannot fetch routing policy")
            return None

        headers: dict[str, str] = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        if self._policy and self._policy.etag:
            headers["If-None-Match"] = self._policy.etag

        try:
            with httpx.Client(timeout=5.0) as client:
                resp = client.get(
                    f"{self.api_base}/route/policy",
                    headers=headers,
                )

            if resp.status_code == 304:
                # Not modified — refresh TTL on existing policy
                if self._policy:
                    self._policy.fetched_at = time.time()
                    self._save_to_disk(self._policy)
                return self._policy

            if resp.status_code != 200:
                logger.debug("Policy fetch returned HTTP %d", resp.status_code)
                return None

            data = resp.json()

            # Extract TTL from Cache-Control header if present
            cc = resp.headers.get("cache-control", "")
            ttl = data.get("ttl_seconds", 3600)
            if "max-age=" in cc:
                try:
                    ttl = int(cc.split("max-age=")[1].split(",")[0].strip())
                except (ValueError, IndexError):
                    pass
            data["ttl_seconds"] = ttl

            etag = resp.headers.get("etag", "")
            data["etag"] = etag
            data["fetched_at"] = time.time()

            return RoutingPolicy.from_dict(data)

        except Exception:
            logger.debug("Failed to fetch routing policy from server", exc_info=True)
            return None

    def _load_from_disk(self) -> Optional[RoutingPolicy]:
        """Load cached policy from disk."""
        path = _CACHE_DIR / _POLICY_FILE
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            return RoutingPolicy.from_dict(data)
        except (FileNotFoundError, json.JSONDecodeError, KeyError):
            return None

    def _save_to_disk(self, policy: RoutingPolicy) -> None:
        """Persist policy to disk cache."""
        path = _CACHE_DIR / _POLICY_FILE
        try:
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(
                json.dumps(policy.to_dict(), indent=2),
                encoding="utf-8",
            )
        except OSError:
            logger.debug("Failed to write policy cache to %s", path, exc_info=True)


# ---------------------------------------------------------------------------
# Deterministic detection (thin — delegates to server when online)
# ---------------------------------------------------------------------------


def check_deterministic(query: str) -> Optional[DeterministicResult]:
    """Check if a query looks like pure arithmetic.

    In online mode the server handles actual evaluation.  This function
    only detects obvious math patterns and returns ``None`` for anything
    that needs a model.  Returns a ``DeterministicResult`` stub when a
    math pattern is detected, suitable for the server to evaluate.
    """
    text = query.strip().rstrip("?.")

    # Strip common prefixes
    for prefix in (
        "what is ",
        "what's ",
        "calculate ",
        "compute ",
        "evaluate ",
        "how much is ",
        "solve ",
        "find ",
    ):
        if text.lower().startswith(prefix):
            text = text[len(prefix) :]
            break

    if not text:
        return None

    # Quick check: if remaining text is purely numeric/operator characters
    # or contains a known math function call, treat as arithmetic
    if _MATH_PATTERN.search(text):
        # Reject if there are alphabetic chars that aren't math function names
        check = text
        for fname in (
            "sqrt",
            "abs",
            "round",
            "ceil",
            "floor",
            "log2",
            "log10",
            "log",
            "sin",
            "cos",
            "tan",
            "pi",
        ):
            check = check.replace(fname, "")
        check = re.sub(r"\bof\b", "", check, flags=re.IGNORECASE)
        if re.search(r"[a-zA-Z_]", check):
            return None

        # Basic safe evaluation for simple expressions
        result = _safe_eval(text)
        if result is not None:
            return DeterministicResult(
                answer=_format_number(result),
                method="arithmetic",
                confidence=1.0,
            )

    return None


def _safe_eval(expr: str) -> Optional[float]:
    """Evaluate a simple arithmetic expression safely via ast."""
    import ast
    import math
    import operator

    # Normalise: ^ to **, implicit multiplication
    expr = re.sub(r"(\d)\s*\^\s*(\d)", r"\1**\2", expr)
    expr = re.sub(r"(\d)\s*\(", r"\1*(", expr)

    # Handle percentage: "N% of M"
    pct = re.match(r"^([\d.]+)\s*%\s*(?:of)\s+([\d.]+)$", expr, re.IGNORECASE)
    if pct:
        return (float(pct.group(1)) / 100.0) * float(pct.group(2))

    _SAFE_NODES: tuple[type[ast.AST], ...] = (
        ast.Expression,
        ast.Constant,
        ast.UnaryOp,
        ast.UAdd,
        ast.USub,
        ast.BinOp,
        ast.Add,
        ast.Sub,
        ast.Mult,
        ast.Div,
        ast.FloorDiv,
        ast.Mod,
        ast.Pow,
        ast.Call,
        ast.Name,
        ast.Load,
    )
    if hasattr(ast, "Num"):
        _SAFE_NODES = _SAFE_NODES + (ast.Num,)

    _FUNCS: dict[str, Any] = {
        "sqrt": math.sqrt,
        "abs": abs,
        "round": round,
        "ceil": math.ceil,
        "floor": math.floor,
        "log": math.log,
        "log2": math.log2,
        "log10": math.log10,
        "sin": math.sin,
        "cos": math.cos,
        "tan": math.tan,
    }
    _CONSTS: dict[str, float] = {"pi": math.pi, "e": math.e}

    def _eval_node(node: ast.AST) -> float:
        if isinstance(node, ast.Expression):
            return _eval_node(node.body)
        if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
            return float(node.value)
        if hasattr(ast, "Num") and isinstance(node, ast.Num):
            n = node.n  # type: ignore[attr-defined]
            if isinstance(n, (int, float)):
                return float(n)
            raise ValueError
        if isinstance(node, ast.UnaryOp):
            val = _eval_node(node.operand)
            if isinstance(node.op, ast.UAdd):
                return +val
            if isinstance(node.op, ast.USub):
                return -val
            raise ValueError
        if isinstance(node, ast.BinOp):
            left, right = _eval_node(node.left), _eval_node(node.right)
            ops = {
                ast.Add: operator.add,
                ast.Sub: operator.sub,
                ast.Mult: operator.mul,
                ast.Div: operator.truediv,
                ast.FloorDiv: operator.floordiv,
                ast.Mod: operator.mod,
                ast.Pow: operator.pow,
            }
            fn = ops.get(type(node.op))
            if fn is None:
                raise ValueError
            if isinstance(node.op, ast.Pow) and abs(right) > 10000:
                raise ValueError("Exponent too large")
            binop_result: float = fn(left, right)
            return binop_result
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
            if node.func.id not in _FUNCS or node.keywords:
                raise ValueError
            fn_result: float = _FUNCS[node.func.id](*[_eval_node(a) for a in node.args])
            return fn_result
        if isinstance(node, ast.Name) and node.id in _CONSTS:
            return _CONSTS[node.id]
        raise ValueError

    try:
        tree = ast.parse(expr, mode="eval")
    except SyntaxError:
        return None

    for node in ast.walk(tree):
        if not isinstance(node, _SAFE_NODES):
            return None

    try:
        return _eval_node(tree)
    except (ValueError, TypeError, ZeroDivisionError, OverflowError, ArithmeticError):
        return None


def _format_number(value: float) -> str:
    """Format a numeric result for display."""
    import math as _math

    if value == float("inf") or value == float("-inf") or _math.isnan(value):
        return str(value)
    if value == int(value) and abs(value) < 1e15:
        return str(int(value))
    return f"{value:.10g}"


# ---------------------------------------------------------------------------
# QueryRouter
# ---------------------------------------------------------------------------


class QueryRouter:
    """Routes queries to the best model based on server-side policy.

    Parameters
    ----------
    models:
        Dict mapping model name to ``ModelInfo``.
    strategy:
        Routing strategy. Currently only ``"complexity"`` is supported.
    thresholds:
        Two-element tuple ``(low, high)`` defining tier boundaries.
    enable_deterministic:
        If ``True``, check for tier-0 deterministic answers.
    api_base:
        Server URL for policy fetching and server-side routing.
    api_key:
        API key for server authentication.
    """

    def __init__(
        self,
        models: dict[str, ModelInfo],
        strategy: str = "complexity",
        thresholds: tuple[float, float] = (0.3, 0.7),
        enable_deterministic: bool = True,
        api_base: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> None:
        if not models:
            raise OctomilError(code=OctomilErrorCode.INVALID_INPUT, message="At least one model must be provided")
        if strategy != "complexity":
            raise OctomilError(
                code=OctomilErrorCode.INVALID_INPUT,
                message=f"Unknown routing strategy '{strategy}'. Supported strategies: complexity",
            )

        self.models = models
        self.strategy = strategy
        self.thresholds = thresholds
        self.enable_deterministic = enable_deterministic

        self._tier_models: dict[str, str] = {}
        for name, info in models.items():
            if info.tier not in self._tier_models:
                self._tier_models[info.tier] = name

        self._ordered_models: list[str] = sorted(
            models.keys(),
            key=lambda n: models[n].tier_index,
        )

        _base: str = api_base or os.environ.get("OCTOMIL_API_BASE") or "https://api.octomil.com/api/v1"
        _key: str = api_key or os.environ.get("OCTOMIL_API_KEY") or ""
        self._policy_client = PolicyClient(api_base=_base, api_key=_key)

        logger.info(
            "QueryRouter initialised: strategy=%s, models=%s, thresholds=%s, deterministic=%s",
            strategy,
            list(models.keys()),
            thresholds,
            enable_deterministic,
        )

    def route(self, messages: list[dict[str, str]]) -> RoutingDecision:
        """Determine which model should handle this request.

        Parameters
        ----------
        messages:
            OpenAI-style message list (dicts with ``role`` and ``content``).

        Returns
        -------
        RoutingDecision with the chosen model name, complexity score,
        tier, and fallback chain.
        """
        user_text = ""
        for msg in messages:
            if msg.get("role") == "user":
                user_text = msg.get("content", "")

        # Tier 0: deterministic
        if self.enable_deterministic and user_text:
            det = check_deterministic(user_text)
            if det is not None:
                logger.debug(
                    "Deterministic hit: %r -> %s (method=%s)",
                    user_text,
                    det.answer,
                    det.method,
                )
                return RoutingDecision(
                    model_name="",
                    complexity_score=0.0,
                    tier="deterministic",
                    strategy=self.strategy,
                    deterministic_result=det,
                )

        # Apply policy-based routing
        policy = self._policy_client.get_policy()
        target_tier, score = self._apply_policy(user_text, policy)

        model_name = self._resolve_model(target_tier)
        fallback_chain = self._build_fallback_chain(model_name)

        decision = RoutingDecision(
            model_name=model_name,
            complexity_score=round(score, 4),
            tier=target_tier,
            strategy=self.strategy,
            fallback_chain=fallback_chain,
        )

        logger.debug(
            "Routing decision: score=%.3f, tier=%s, model=%s",
            score,
            target_tier,
            model_name,
        )

        return decision

    def _apply_policy(self, query: str, policy: RoutingPolicy) -> tuple[str, float]:
        """Apply cached policy to classify query into a tier.

        Uses the iOS-aligned scoring formula: linear interpolation of word
        count between thresholds combined with a complex-indicator fraction,
        weighted by ``length_weight`` and ``indicator_weight``.

        Returns (tier, complexity_score).
        """
        word_count = len(query.split())
        query_lower = query.lower()

        # Word score: linear interpolation between thresholds (0.0-1.0)
        if word_count <= policy.fast_max_words:
            word_score = 0.0
        elif word_count >= policy.quality_min_words:
            word_score = 1.0
        else:
            span = policy.quality_min_words - policy.fast_max_words
            word_score = (word_count - policy.fast_max_words) / span if span > 0 else 0.5

        # Indicator score: fraction of complex indicators found
        indicator_count = sum(1 for kw in policy.complex_indicators if kw in query_lower)
        normalizor = policy.indicator_normalizor if policy.indicator_normalizor > 0 else 1.0
        indicator_score = min(indicator_count / normalizor, 1.0)

        # Combined score
        score = min(
            word_score * policy.length_weight + indicator_score * policy.indicator_weight,
            1.0,
        )

        # Tier assignment
        if score < policy.fast_threshold:
            tier = "fast"
        elif score >= policy.quality_threshold:
            tier = "quality"
        else:
            tier = "balanced"

        return tier, round(score, 4)

    def _resolve_model(self, target_tier: str) -> str:
        """Find the best available model for a tier, falling back upward."""
        if target_tier in self._tier_models:
            return self._tier_models[target_tier]

        try:
            target_idx = TIER_ORDER.index(target_tier)
        except ValueError:
            target_idx = 0

        for tier in TIER_ORDER[target_idx:]:
            if tier in self._tier_models:
                return self._tier_models[tier]

        for tier in reversed(TIER_ORDER[:target_idx]):
            if tier in self._tier_models:
                return self._tier_models[tier]

        return self._ordered_models[0]

    def _build_fallback_chain(self, primary: str) -> list[str]:
        """Build ordered list of fallback models (excluding primary)."""
        primary_idx = self.models[primary].tier_index
        higher = [n for n in self._ordered_models if n != primary and self.models[n].tier_index > primary_idx]
        lower = [n for n in self._ordered_models if n != primary and self.models[n].tier_index < primary_idx]
        same = [n for n in self._ordered_models if n != primary and self.models[n].tier_index == primary_idx]
        return higher + same + lower

    def route_decomposed(
        self,
        messages: list[dict[str, str]],
    ) -> Union[RoutingDecision, DecomposedRoutingDecision]:
        """Route with optional query decomposition.

        If the query contains multiple independent sub-tasks, returns a
        ``DecomposedRoutingDecision`` with per-task routing.  Otherwise
        returns a standard ``RoutingDecision``.
        """
        from .decomposer import QueryDecomposer

        decomposer = QueryDecomposer()
        decomposition = decomposer.decompose(messages)

        if not decomposition.decomposed:
            return self.route(messages)

        sub_decisions: list[RoutingDecision] = []
        for task in decomposition.tasks:
            sub_messages = [{"role": "user", "content": task.text}]
            for msg in messages:
                if msg.get("role") == "system":
                    sub_messages.insert(0, msg)
                    break
            decision = self.route(sub_messages)
            sub_decisions.append(decision)

        return DecomposedRoutingDecision(
            sub_decisions=sub_decisions,
            tasks=decomposition.tasks,
            original_query=decomposition.original_query,
        )

    def get_fallback(self, failed_model: str) -> Optional[str]:
        """Return the next model to try after ``failed_model`` fails.

        Returns ``None`` if no fallback is available.
        """
        failed_idx = self.models.get(failed_model, ModelInfo(name="")).tier_index

        for name in self._ordered_models:
            if name != failed_model and self.models[name].tier_index > failed_idx:
                return name
        for name in self._ordered_models:
            if name != failed_model and self.models[name].tier_index == failed_idx:
                return name
        for name in reversed(self._ordered_models):
            if name != failed_model and self.models[name].tier_index < failed_idx:
                return name
        return None


def assign_tiers(
    model_names: list[str],
    thresholds: tuple[float, float] = (0.3, 0.7),
) -> dict[str, ModelInfo]:
    """Auto-assign tiers to an ordered list of model names.

    Assumes ``model_names`` is ordered from smallest to largest.
    """
    n = len(model_names)
    if n == 0:
        return {}
    if n == 1:
        return {model_names[0]: ModelInfo(name=model_names[0], tier="balanced")}
    if n == 2:
        return {
            model_names[0]: ModelInfo(name=model_names[0], tier="fast"),
            model_names[1]: ModelInfo(name=model_names[1], tier="quality"),
        }

    result: dict[str, ModelInfo] = {}
    fast_end = n // 3
    quality_start = n - (n // 3)

    for i, name in enumerate(model_names):
        if i < fast_end:
            tier = "fast"
        elif i >= quality_start:
            tier = "quality"
        else:
            tier = "balanced"
        result[name] = ModelInfo(name=name, tier=tier)

    return result
