"""LLM Router —— 统一模型路由层（#875 · 影子模式先行）

设计基线（leagent_accel_belt_design.md §C）：
- 统一 Web Claw 与 IM Bot Claw 的模型选择，消除两端重复的 failover / 主密钥硬编码。
- 决策顺序：显式指定 → 租户配额 → 能力过滤 → 成本分层 → 健康 failover（熔断）。
- 灰度策略：**默认影子模式**（`LLM_ROUTER_SHADOW=1`）—— Router 计算结果只记日志 +
  `claw_type` 标签 metrics，实际仍走旧逻辑；一致率 ≥98% 后把开关置 0 切主路径。
- 上下文三元组（tenant_id + claw_type + trace_id）经 `contextvars` 注入，零签名侵入
  （复用 observability.py 的 trace_id 同款机制）。

对外 API：
    set_route_ctx(tenant_id="", claw_type="web")   # 在入口注入上下文
    get_route_ctx() -> dict                        # {tenant_id, claw_type, trace_id}
    route(capability, db, *, prefer_model_id=None, legacy_primary=None, shadow=None)
        -> RouteDecision(models: list[AIModel], primary, reason, matched_legacy, ...)
    resolve_api_key(api_url, model_api_key, provider_api_key) -> str|None
        # 收敛 LiteLLM 主密钥硬编码：litellm 端点统一返回 LITELLM_MASTER_KEY
    report_failure(model_id, claw_type="web") / report_success(model_id)
        # 健康熔断：Redis 记录连续失败（TTL 60s），无 Redis 时降级内存
"""

from __future__ import annotations

import os
import time
import logging
import threading
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# 影子模式开关：默认开启 → 只观测，不改流量
# ---------------------------------------------------------------------------
SHADOW = os.getenv("LLM_ROUTER_SHADOW", "1") not in ("0", "false", "False", "")

_LITELLM_MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY") or "sk-litellm-master-key-aiflow-2026"
_LITELLM_HOST_MARK = "qihu-aiflow-litellm"

# 媒体能力（含 "text" 但本质是媒体模型，不能当文本 failover 备选）
_MEDIA_CAPS = {"image", "video", "audio", "asr", "embedding", "rerank", "ocr"}


# ---------------------------------------------------------------------------
# 上下文三元组：tenant_id + claw_type（trace_id 复用 observability）
# ---------------------------------------------------------------------------
_route_tenant_var: ContextVar[str] = ContextVar("aiflow_route_tenant", default="")
_route_claw_var: ContextVar[str] = ContextVar("aiflow_route_claw", default="web")


def set_route_ctx(*, tenant_id: str = "", claw_type: str = "web") -> None:
    """在请求入口注入路由上下文。tenant_id/claw_type 为空时保留默认值。"""
    if tenant_id:
        _route_tenant_var.set(tenant_id)
    _route_claw_var.set(claw_type or "web")


def get_route_ctx() -> dict:
    try:
        from app.core.observability import get_trace_id
        _trace = get_trace_id()
    except Exception:
        _trace = "-"
    return {
        "tenant_id": _route_tenant_var.get(),
        "claw_type": _route_claw_var.get(),
        "trace_id": _trace,
    }


# ---------------------------------------------------------------------------
# 决策结果
# ---------------------------------------------------------------------------
@dataclass
class RouteDecision:
    capability: str
    claw_type: str
    tenant_id: str
    models: list = field(default_factory=list)      # 有序候选链 list[AIModel]
    primary: object | None = None                   # 首选 AIModel
    reason: str = ""
    shadow: bool = True
    matched_legacy: Optional[bool] = None           # 与旧逻辑首选是否一致


# ---------------------------------------------------------------------------
# 能力过滤：复用 capability_registry + 内置 _is_text_only
# ---------------------------------------------------------------------------
def _is_text_only(m) -> bool:
    """纯文本模型才能作为文本 failover 备选（消除混入图像/语音模型导致空内容的坑）。"""
    caps = m.capabilities or []
    return "text" in caps and not any(c in _MEDIA_CAPS for c in caps)


def _fallbacks(db, exclude_ids: set, limit: int = 4) -> list:
    """对齐 IM Bot 的备选构造：全部启用、按 is_default 降序、仅纯文本、取前 limit。"""
    from app.models import AIModel
    rows = (
        db.query(AIModel)
        .filter(AIModel.is_enabled == True)  # noqa: E712
        .order_by(AIModel.is_default.desc())
        .all()
    )
    out = [m for m in rows if m.id not in exclude_ids and _is_text_only(m)]
    return out[:limit]


def _tenant_authorized(m, tenant_id: str) -> bool:
    """模型归属校验：全局模型(无 tenant)任何人可用；否则仅所属租户可用。"""
    mt = getattr(m, "tenant_id", None)
    if not mt:
        return True
    if not tenant_id:
        return False
    return mt == tenant_id


# ---------------------------------------------------------------------------
# 健康熔断：Redis 连续失败计数（TTL 60s），无 Redis 降级内存
# ---------------------------------------------------------------------------
_BREAKER_TTL = 60.0
_BREAKER_LOCAL: dict = {}
_BREAKER_LOCK = threading.Lock()


def _redis():
    try:
        import redis
        from app.core.config import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD
        return redis.Redis(
            host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD,
            socket_connect_timeout=1, socket_timeout=1,
        )
    except Exception:
        return None


def report_failure(model_id: str, claw_type: str = "web") -> None:
    if not model_id:
        return
    r = _redis()
    key = f"aiflow:router:fail:{model_id}"
    if r:
        try:
            r.incr(key)
            r.expire(key, int(_BREAKER_TTL))
            return
        except Exception:
            pass
    with _BREAKER_LOCK:
        e = _BREAKER_LOCAL.setdefault(model_id, {"fails": 0, "until": 0.0})
        e["fails"] += 1
        e["until"] = time.time() + _BREAKER_TTL


def report_success(model_id: str) -> None:
    if not model_id:
        return
    r = _redis()
    key = f"aiflow:router:fail:{model_id}"
    if r:
        try:
            r.delete(key)
            return
        except Exception:
            pass
    with _BREAKER_LOCK:
        _BREAKER_LOCAL.pop(model_id, None)


def _fail_count(model_id: str) -> int:
    if not model_id:
        return 0
    r = _redis()
    key = f"aiflow:router:fail:{model_id}"
    if r:
        try:
            v = r.get(key)
            return int(v) if v else 0
        except Exception:
            pass
    with _BREAKER_LOCK:
        e = _BREAKER_LOCAL.get(model_id)
        if e and e["until"] > time.time():
            return e["fails"]
        _BREAKER_LOCAL.pop(model_id, None)
        return 0


def _apply_breaker(models: list, claw_type: str) -> list:
    """失败次数少的模型优先；0 失败保持原序。"""
    if not models:
        return models
    scored = [( _fail_count(getattr(m, "id", "")), i, m) for i, m in enumerate(models)]
    scored.sort(key=lambda x: (x[0], x[1]))
    return [m for _, _, m in scored]


# ---------------------------------------------------------------------------
# 主密钥收敛：litellm 端点统一用 master key
# ---------------------------------------------------------------------------
def resolve_api_key(api_url: Optional[str], model_api_key: Optional[str],
                    provider_api_key: Optional[str]) -> Optional[str]:
    if api_url and _LITELLM_HOST_MARK in api_url:
        return _LITELLM_MASTER_KEY
    return model_api_key or provider_api_key


# ---------------------------------------------------------------------------
# metrics：复用 otel meter，带 claw_type 标签（失败安全）
# ---------------------------------------------------------------------------
_counter = None
_match_counter = None
_hist = None


def _metric_objs():
    global _counter, _match_counter, _hist
    if _counter is not None:
        return _counter, _match_counter, _hist
    try:
        from app.core.otel import get_meter
        m = get_meter()
        if m is None:
            return None, None, None
        _counter = m.create_counter("aiflow_llm_route_total", "LLM router decisions (shadow+live)")
        _match_counter = m.create_counter(
            "aiflow_llm_route_legacy_match_total", "Router vs legacy primary consistency"
        )
        _hist = m.create_histogram("aiflow_llm_route_ms", "Router decision latency ms")
    except Exception:
        return None, None, None
    return _counter, _match_counter, _hist


def _emit_metrics(capability: str, claw_type: str, primary, shadow: bool, elapsed_ms: float) -> None:
    c, mc, h = _metric_objs()
    if c is None:
        return
    try:
        c.add(1, {
            "claw_type": claw_type,
            "capability": capability,
            "decision": primary.name if primary else "none",
            "shadow": str(shadow),
        })
        if h:
            h.record(elapsed_ms, {"claw_type": claw_type})
    except Exception:
        pass


def _emit_match(claw_type: str, matched: bool) -> None:
    c, mc, h = _metric_objs()
    if mc is None:
        return
    try:
        mc.add(1, {"claw_type": claw_type, "matched": "true" if matched else "false"})
    except Exception:
        pass


# ---------------------------------------------------------------------------
# route：同步（与 resolve_auto_model 一致，可直接在 async 引擎内调用）
# ---------------------------------------------------------------------------
def route(capability: str, db, *, prefer_model_id: Optional[str] = None,
          legacy_primary: Optional[str] = None, shadow: Optional[bool] = None) -> RouteDecision:
    _t0 = time.perf_counter()
    shadow = SHADOW if shadow is None else shadow
    ctx = get_route_ctx()
    claw_type = ctx["claw_type"]
    tenant_id = ctx["tenant_id"]
    reasons: list = []

    models: list = []
    try:
        from app.models import AIModel
        from app.capability_registry import resolve_auto_model

        # 1) 显式指定
        if prefer_model_id:
            m = (
                db.query(AIModel)
                .filter(AIModel.id == prefer_model_id, AIModel.is_enabled == True)  # noqa: E712
                .first()
            )
            if m and _tenant_authorized(m, tenant_id):
                models = [m] + _fallbacks(db, exclude_ids={m.id})
                reasons.append(f"explicit={m.name}")
            else:
                reasons.append("explicit rejected (disabled or tenant mismatch)")

        # 2) 能力过滤：平台默认 + 文本备选（配额/成本分档为预留钩子，未强校验）
        if not models:
            primary = resolve_auto_model(capability, db)
            if primary and primary.is_enabled:
                models = [primary] + _fallbacks(db, exclude_ids={primary.id})
                reasons.append(f"auto={primary.name}")
            else:
                reasons.append("no primary resolved")

        # 5) 健康 failover 重排
        models = _apply_breaker(models, claw_type)
    except Exception as _e:
        logger.warning(f"[llm_router] route error cap={capability} claw={claw_type}: {_e}")
        reasons.append(f"error:{type(_e).__name__}")

    primary = models[0] if models else None
    matched = None
    if legacy_primary is not None:
        matched = bool(primary and primary.name == legacy_primary)

    elapsed_ms = (time.perf_counter() - _t0) * 1000
    _emit_metrics(capability, claw_type, primary, shadow, elapsed_ms)
    if matched is not None:
        _emit_match(claw_type, matched)

    decision = RouteDecision(
        capability=capability,
        claw_type=claw_type,
        tenant_id=tenant_id,
        models=models,
        primary=primary,
        reason="; ".join(reasons) or "default",
        shadow=shadow,
        matched_legacy=matched,
    )
    logger.info(
        f"[llm_router] claw={claw_type} cap={capability} primary="
        f"{primary.name if primary else None} shadow={shadow} matched={matched} "
        f"reason={decision.reason}"
    )
    return decision


# 预留钩子（#875 后续迭代启用）：租户配额 / 成本分档
# def _check_tenant_quota(tenant_id, capability) -> tuple[bool, str]: ...
# def _cost_tier_models(tenant_id, capability) -> list: ...