"""Lane orchestration A/B/C with Redis session context."""

from __future__ import annotations

import hashlib
import json
import logging
import os
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from rules_loader import load_model_matrix, load_orchestration, load_routing_rules

log = logging.getLogger("orchestrator")

LANE_PREFIX = {"A": "a", "B": "b", "C": "c"}
TIER_SUFFIX = {
    "SIMPLE": "simple",
    "MEDIUM_OPS": "medium-ops",
    "MEDIUM_CODE": "medium-code",
    "COMPLEX": "complex",
    "REASONING": "reasoning",
    "VISION_OCR": "vision-ocr",
    "VISION_UI": "vision",
}


class Tier(str, Enum):
    SIMPLE = "SIMPLE"
    MEDIUM_OPS = "MEDIUM_OPS"
    MEDIUM_CODE = "MEDIUM_CODE"
    COMPLEX = "COMPLEX"
    REASONING = "REASONING"
    VISION_OCR = "VISION_OCR"
    VISION_UI = "VISION_UI"


@dataclass
class SessionCtx:
    escalation_level: int = 0
    last_prompt_hash: str = ""
    turn_count: int = 0
    budget_pct: float = 0.0


@dataclass
class RouteDecision:
    tier: Tier
    lane: str
    model: str
    quality_mode: str
    escalation_level: int
    confidence: float
    delegated_internal: bool = False


class SessionStore:
    def __init__(self, redis_url: str | None) -> None:
        self._redis = None
        if redis_url:
            try:
                import redis

                self._redis = redis.from_url(redis_url, decode_responses=True)
                self._redis.ping()
            except Exception as exc:  # noqa: BLE001
                log.warning("Redis unavailable (%s), using in-memory sessions", exc)
        self._memory: dict[str, SessionCtx] = {}
        orch = load_orchestration()
        self._prefix = orch.get("redis", {}).get("key_prefix", "ai-router:session:")
        self._ttl = int(orch.get("redis", {}).get("ttl_sec", 1800))

    def get(self, session_id: str) -> SessionCtx:
        if not session_id:
            return SessionCtx()
        key = f"{self._prefix}{session_id}"
        if self._redis:
            raw = self._redis.get(key)
            if raw:
                data = json.loads(raw)
                return SessionCtx(**data)
            return SessionCtx()
        return self._memory.get(session_id, SessionCtx())

    def save(self, session_id: str, ctx: SessionCtx) -> None:
        if not session_id:
            return
        key = f"{self._prefix}{session_id}"
        payload = json.dumps(
            {
                "escalation_level": ctx.escalation_level,
                "last_prompt_hash": ctx.last_prompt_hash,
                "turn_count": ctx.turn_count,
                "budget_pct": ctx.budget_pct,
            }
        )
        if self._redis:
            self._redis.setex(key, self._ttl, payload)
        else:
            self._memory[session_id] = ctx


class Classifier:
    def __init__(self) -> None:
        rules = load_routing_rules()
        gw = rules.get("gateway", {})
        self._simple_re = re.compile(
            gw.get("simple_patterns", [r"^(hello)\\b"])[0],
            re.IGNORECASE,
        )
        self._complex = self._kw_re(gw.get("complex_keywords", []))
        self._reasoning = self._kw_re(gw.get("reasoning_keywords", []))
        self._medium_ops = self._kw_re(gw.get("medium_ops_keywords", []))
        self._ocr = self._kw_re(gw.get("ocr_keywords", []))
        self._escalation = self._kw_re(gw.get("escalation_keywords", []))
        wt = gw.get("word_thresholds", {})
        self._simple_max = int(wt.get("simple_max_words", 12))
        self._simple_q_max = int(wt.get("simple_question_max_words", 25))
        self._complex_min = int(wt.get("complex_min_words", 400))
        self._low_conf = float(rules.get("gateway", {}).get("confidence", {}).get("low_threshold", 0.6))
        self._code_block = re.compile(r"```[\s\S]*?```|`[^`]+`")

    @staticmethod
    def _kw_re(keywords: list[str]) -> re.Pattern[str]:
        if not keywords:
            return re.compile(r"(?!x)x")
        escaped = [re.escape(k) for k in keywords]
        return re.compile("|".join(escaped), re.IGNORECASE)

    def wants_escalation(self, text: str, prompt_hash: str, prev_hash: str) -> bool:
        if self._escalation.search(text):
            return True
        return bool(prompt_hash and prompt_hash == prev_hash)

    def classify(
        self,
        messages: list[dict[str, Any]],
        *,
        has_image: bool,
        text: str,
    ) -> tuple[Tier, float]:
        if has_image:
            if self._ocr.search(text) or len(text.split()) < 30:
                return Tier.VISION_OCR, 0.95
            return Tier.VISION_UI, 0.9

        text = text.strip()
        if not text:
            return Tier.SIMPLE, 0.9

        words = len(text.split())
        if words <= self._simple_max and self._simple_re.search(text):
            return Tier.SIMPLE, 0.92
        if self._reasoning.search(text):
            return Tier.REASONING, 0.88
        if self._complex.search(text):
            return Tier.COMPLEX, 0.88
        if self._code_block.search(text):
            return Tier.MEDIUM_CODE, 0.85
        if self._medium_ops.search(text):
            return Tier.MEDIUM_OPS, 0.82
        if words <= self._simple_q_max and "?" in text and not self._code_block.search(text):
            return Tier.SIMPLE, 0.75
        if words >= self._complex_min:
            return Tier.COMPLEX, 0.7
        return Tier.MEDIUM_OPS, 0.55


class Orchestrator:
    def __init__(self, session_store: SessionStore) -> None:
        self.sessions = session_store
        self.classifier = Classifier()
        self._orch = load_orchestration()
        self._matrix = load_model_matrix()
        self._models: dict[str, dict] = self._matrix.get("models", {})
        self._no_esc = {
            name for name, cfg in self._models.items() if cfg.get("no_escalation")
        }
        self.default_quality = os.environ.get(
            "DEFAULT_QUALITY_MODE",
            self._orch.get("default_quality_mode", "auto"),
        )

    def _lane_from_mode(self, quality_mode: str) -> str | None:
        if quality_mode == "auto":
            return None
        return self._orch.get("quality_mode_map", {}).get(quality_mode)

    def _start_lane(self, tier: Tier) -> str:
        return self._orch.get("start_lanes", {}).get(tier.value, "A")

    def _bump_lane(self, lane: str, levels: int) -> str:
        order: list[str] = self._orch.get("lane_order", ["A", "B", "C"])
        try:
            idx = order.index(lane)
        except ValueError:
            return lane
        return order[min(idx + levels, len(order) - 1)]

    def _cap_lane(self, lane: str, budget_pct: float) -> str:
        caps = self._orch.get("budget_caps", {})
        hard = float(caps.get("hard_pct", 95))
        warn = float(caps.get("warn_pct", 80))
        if budget_pct >= hard:
            return caps.get("hard_max_lane", "A")
        if budget_pct >= warn:
            max_lane = caps.get("warn_max_lane", "B")
            order = self._orch.get("lane_order", ["A", "B", "C"])
            if order.index(lane) > order.index(max_lane):
                return max_lane
        return lane

    def _model_name(self, lane: str, tier: Tier) -> str:
        prefix = LANE_PREFIX.get(lane, "a")
        suffix = TIER_SUFFIX[tier]
        return f"{prefix}-{suffix}"

    def resolve(
        self,
        messages: list[dict[str, Any]],
        *,
        quality_mode: str | None = None,
        session_id: str = "",
        text: str = "",
        has_image: bool = False,
        token_estimate: int = 0,
    ) -> RouteDecision:
        mode = quality_mode or self.default_quality
        tier, confidence = self.classifier.classify(messages, has_image=has_image, text=text)
        ctx = self.sessions.get(session_id)
        prompt_hash = hashlib.sha256(text.encode()).hexdigest()[:16]

        if mode != "auto":
            lane = self._lane_from_mode(mode) or "A"
            escalation = 0
        else:
            lane = self._start_lane(tier)
            escalation = ctx.escalation_level
            if tier != Tier.VISION_OCR and self.classifier.wants_escalation(
                text, prompt_hash, ctx.last_prompt_hash
            ):
                escalation += 1
            min_tokens = int(
                self._orch.get("escalation", {}).get("context_tokens_min_lane_b", 32000)
            )
            if token_estimate > min_tokens:
                order = self._orch.get("lane_order", ["A", "B", "C"])
                if order.index(lane) < order.index("B"):
                    lane = "B"
            lane = self._bump_lane(lane, escalation)
            lane = self._cap_lane(lane, ctx.budget_pct)

        model = self._model_name(lane, tier)
        if model not in self._models:
            log.warning("model %s missing from matrix, fallback smart-router-internal", model)
            return RouteDecision(
                tier=tier,
                lane=lane,
                model="smart-router-internal",
                quality_mode=mode,
                escalation_level=escalation,
                confidence=confidence,
                delegated_internal=True,
            )

        delegated = confidence < self.classifier._low_conf and tier in (
            Tier.MEDIUM_OPS,
            Tier.SIMPLE,
        )

        return RouteDecision(
            tier=tier,
            lane=lane,
            model="smart-router-internal" if delegated else model,
            quality_mode=mode,
            escalation_level=escalation,
            confidence=confidence,
            delegated_internal=delegated,
        )

    def after_request(
        self,
        session_id: str,
        *,
        prompt_hash: str,
        success: bool,
        escalate: bool,
    ) -> None:
        if not session_id:
            return
        ctx = self.sessions.get(session_id)
        ctx.last_prompt_hash = prompt_hash
        ctx.turn_count += 1
        if success and not escalate:
            ctx.escalation_level = 0
        elif escalate:
            ctx.escalation_level = min(ctx.escalation_level + 1, 2)
        self.sessions.save(session_id, ctx)