from __future__ import annotations

from contextlib import contextmanager, nullcontext
from typing import Any, Iterator

from niannian.llm.providers.exclusive_lock import (
    ProviderExclusiveBusy,
    exclusive_provider_lock,
    is_exclusive_provider,
    resolve_exclusive_providers,
    resolve_lock_dir,
    resolve_lock_timeout,
)
from niannian.llm.providers.registry import create_llm_provider, is_llm_provider_name
from niannian.llm.spec import parse_provider_model
from niannian.llm.types import (
    CompletionResult,
    EmbeddingResult,
    LLMUnavailable,
    Message,
    OcrResult,
    ProviderSpec,
    RerankResult,
    RouterTrace,
    StreamChunk,
)
from niannian.personas import MODULE_PERSONAS
from niannian.token.context import token_usage_operation, token_usage_user_id
from niannian.token.usage_store import TokenUsageStore

_SLOT_ORDER = ("primary", "module_fallback", "global_fallback")

_MODULE_TEMPERATURE: dict[str, float] = {
    "right": 0.7,
    "left": 0.2,
    "brainstem": 0.0,
    "hippocampus": 0.3,
    "cerebellum": 0.0,
}

# 小织 embed / rerank 只用本模块 slot，不继承 api.global_fallback（避免 Chat 模型误填）
_MODULES_EXCLUDE_API_GLOBAL_FALLBACK = frozenset({"cerebellum", "retrieval_rerank", "ocr"})


class LLMRouter:
    """按脑区三层 slot 路由云端 LLM：primary → module_fallback → global_fallback。"""

    def __init__(
        self,
        config: dict[str, Any],
        providers_config: dict[str, Any] | None = None,
        *,
        show_trace: bool = True,
    ) -> None:
        self._config = config
        self._api_cfg = config.get("api", {})
        self._providers_cfg = providers_config or {}
        self._show_trace = show_trace
        self._provider_cache: dict[str, LLMProvider] = {}
        self._exclusive_providers = resolve_exclusive_providers(self._api_cfg)
        self._usage_store: TokenUsageStore | None = None

    def set_usage_store(self, store: TokenUsageStore | None) -> None:
        self._usage_store = store

    def reload_llm(self, config: dict[str, Any], providers_config: dict[str, Any]) -> None:
        self._config = config
        self._api_cfg = config.get("api", {})
        self._providers_cfg = providers_config
        self._exclusive_providers = resolve_exclusive_providers(self._api_cfg)
        self._provider_cache.clear()

    def is_module_configured(self, module: str) -> bool:
        return bool(self._slot_chain(module))

    def complete(
        self,
        module: str,
        messages: list[Message],
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        json_mode: bool = False,
        trace: RouterTrace | None = None,
    ) -> CompletionResult:
        chain = self._available_chain(module)
        if not chain:
            raise LLMUnavailable(module, "未配置 LLM slot 或无可用 API Key")

        errors: list[str] = []
        for spec in chain:
            try:
                with self._provider_session(module, spec, trace) as (provider, label):
                    if provider is None or not provider.is_available():
                        errors.append(f"{spec.slot}={spec.provider}:{spec.model} 无 Key")
                        continue
                    temp = temperature if temperature is not None else _MODULE_TEMPERATURE.get(module, 0.7)
                    self._trace(trace, module, f"调用 {label} ({spec.slot})")
                    result = provider.complete(
                        spec.model,
                        messages,
                        temperature=temp,
                        max_tokens=max_tokens,
                        json_mode=json_mode,
                    )
                    result.slot = spec.slot
                    self._trace(trace, module, f"完成 {label}")
                    self._record_usage(module, spec.provider, spec.model, result.usage, operation="complete")
                    return result
            except ProviderExclusiveBusy as exc:
                errors.append(str(exc))
                self._trace(trace, module, f"智谱占用中，降级: {spec.provider}:{spec.model}")
            except Exception as exc:
                msg = f"{spec.slot} {spec.provider}:{spec.model} 失败: {exc}"
                errors.append(msg)
                self._trace(trace, module, f"降级: {msg}")

        raise LLMUnavailable(module, attempts=errors)

    def stream_complete(
        self,
        module: str,
        messages: list[Message],
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        trace: RouterTrace | None = None,
    ) -> Iterator[StreamChunk]:
        chain = self._available_chain(module)
        if not chain:
            raise LLMUnavailable(module, "未配置 LLM slot 或无可用 API Key")

        errors: list[str] = []
        for spec in chain:
            try:
                with self._provider_session(module, spec, trace) as (provider, label):
                    if provider is None or not provider.is_available():
                        errors.append(f"{spec.slot}={spec.provider}:{spec.model} 无 Key")
                        continue
                    temp = temperature if temperature is not None else _MODULE_TEMPERATURE.get(module, 0.7)
                    self._trace(trace, module, f"流式调用 {label} ({spec.slot})")
                    last_usage = None
                    for chunk in provider.stream_complete(
                        spec.model,
                        messages,
                        temperature=temp,
                        max_tokens=max_tokens,
                    ):
                        if chunk.usage:
                            last_usage = chunk.usage
                        yield chunk
                    self._trace(trace, module, f"流式完成 {label}")
                    self._record_usage(module, spec.provider, spec.model, last_usage, operation="stream")
                    return
            except ProviderExclusiveBusy as exc:
                errors.append(str(exc))
                self._trace(trace, module, f"智谱占用中，降级: {spec.provider}:{spec.model}")
            except Exception as exc:
                msg = f"{spec.slot} {spec.provider}:{spec.model} 失败: {exc}"
                errors.append(msg)
                self._trace(trace, module, f"降级: {msg}")

        raise LLMUnavailable(module, attempts=errors)

    def embed(
        self,
        module: str,
        texts: list[str],
        *,
        trace: RouterTrace | None = None,
    ) -> EmbeddingResult:
        chain = self._available_chain(module)
        if not chain:
            raise LLMUnavailable(module, "未配置 embedding slot 或无可用 API Key")

        errors: list[str] = []
        for spec in chain:
            try:
                with self._provider_session(module, spec, trace) as (provider, label):
                    if provider is None or not provider.is_available():
                        errors.append(f"{spec.slot}={spec.provider}:{spec.model} 无 Key")
                        continue
                    self._trace(trace, module, f"embed {label} ({spec.slot})")
                    result = provider.embed(spec.model, texts)
                    result.slot = spec.slot
                    est = sum(max(1, len(t) // 4) for t in texts)
                    self._record_usage(
                        module,
                        spec.provider,
                        spec.model,
                        None,
                        operation="embed",
                        prompt_tokens=est,
                    )
                    return result
            except ProviderExclusiveBusy as exc:
                errors.append(str(exc))
                self._trace(trace, module, f"智谱占用中，降级: {spec.provider}:{spec.model}")
            except NotImplementedError:
                errors.append(f"{spec.provider} 不支持 embed")
            except Exception as exc:
                msg = f"{spec.slot} {spec.provider}:{spec.model} 失败: {exc}"
                errors.append(msg)
                self._trace(trace, module, f"降级: {msg}")

        raise LLMUnavailable(module, attempts=errors)

    def rerank(
        self,
        module: str,
        query: str,
        documents: list[str],
        *,
        top_n: int | None = None,
        trace: RouterTrace | None = None,
    ) -> RerankResult:
        chain = self._available_chain(module)
        if not chain:
            raise LLMUnavailable(module, "未配置 rerank slot 或无可用 API Key")
        if not documents:
            return RerankResult(scores=[], indices=[], model="", provider="", slot="primary")

        errors: list[str] = []
        for spec in chain:
            try:
                with self._provider_session(module, spec, trace) as (provider, label):
                    if provider is None or not provider.is_available():
                        errors.append(f"{spec.slot}={spec.provider}:{spec.model} 无 Key")
                        continue
                    self._trace(trace, module, f"rerank {label} ({spec.slot}) n={len(documents)}")
                    result = provider.rerank(spec.model, query, documents, top_n=top_n)
                    result.slot = spec.slot
                    est = max(1, (len(query) + sum(len(d) for d in documents)) // 4)
                    self._record_usage(
                        module,
                        spec.provider,
                        spec.model,
                        None,
                        operation="rerank",
                        prompt_tokens=est,
                    )
                    return result
            except ProviderExclusiveBusy as exc:
                errors.append(str(exc))
                self._trace(trace, module, f"智谱占用中，降级: {spec.provider}:{spec.model}")
            except NotImplementedError:
                errors.append(f"{spec.provider} 不支持 rerank")
            except Exception as exc:
                msg = f"{spec.slot} {spec.provider}:{spec.model} 失败: {exc}"
                errors.append(msg)
                self._trace(trace, module, f"降级: {msg}")

        raise LLMUnavailable(module, attempts=errors)

    def ocr(
        self,
        module: str,
        image_b64: str,
        *,
        mime: str = "image/png",
        prompt: str = "请提取图片中的全部文字，保持段落结构。",
        trace: RouterTrace | None = None,
    ) -> OcrResult:
        chain = self._available_chain(module)
        if not chain:
            raise LLMUnavailable(module, "未配置 ocr slot 或无可用 API Key")

        errors: list[str] = []
        for spec in chain:
            try:
                with self._provider_session(module, spec, trace) as (provider, label):
                    if provider is None or not provider.is_available():
                        errors.append(f"{spec.slot}={spec.provider}:{spec.model} 无 Key")
                        continue
                    self._trace(trace, module, f"ocr {label} ({spec.slot})")
                    result = provider.ocr(
                        spec.model,
                        image_b64,
                        mime=mime,
                        prompt=prompt,
                    )
                    result.slot = spec.slot
                    self._record_usage(module, spec.provider, spec.model, getattr(result, "usage", None), operation="ocr")
                    return result
            except ProviderExclusiveBusy as exc:
                errors.append(str(exc))
                self._trace(trace, module, f"智谱占用中，降级: {spec.provider}:{spec.model}")
            except NotImplementedError:
                errors.append(f"{spec.provider} 不支持 ocr")
            except Exception as exc:
                msg = f"{spec.slot} {spec.provider}:{spec.model} 失败: {exc}"
                errors.append(msg)
                self._trace(trace, module, f"降级: {msg}")

        raise LLMUnavailable(module, attempts=errors)

    @contextmanager
    def _provider_session(self, module: str, spec: ProviderSpec, trace: RouterTrace | None):
        timeout = resolve_lock_timeout(self._api_cfg, spec.provider)
        lock_ctx = (
            exclusive_provider_lock(
                spec.provider,
                resolve_lock_dir(),
                timeout_sec=timeout,
            )
            if is_exclusive_provider(spec.provider, exclusive_names=self._exclusive_providers)
            else nullcontext()
        )
        with lock_ctx:
            provider = self._get_provider(spec.provider)
            label = f"{spec.provider}:{spec.model}"
            yield provider, label

    def _slot_chain(self, module: str) -> list[ProviderSpec]:
        module_cfg = self._api_cfg.get(module, {})
        if not isinstance(module_cfg, dict):
            module_cfg = {}
        specs: list[ProviderSpec] = []
        seen: set[str] = set()

        for slot in _SLOT_ORDER[:2]:
            raw = module_cfg.get(slot)
            spec = parse_provider_model(raw)
            if spec is None:
                continue
            key = f"{spec.provider}:{spec.model}"
            if key in seen:
                continue
            seen.add(key)
            specs.append(ProviderSpec(spec.provider, spec.model, slot))

        global_raw = module_cfg.get("global_fallback")
        if global_raw in (None, "") and module not in _MODULES_EXCLUDE_API_GLOBAL_FALLBACK:
            global_raw = self._api_cfg.get("global_fallback")
        spec = parse_provider_model(global_raw)
        if spec is not None:
            key = f"{spec.provider}:{spec.model}"
            if key not in seen:
                specs.append(ProviderSpec(spec.provider, spec.model, "global_fallback"))

        return specs

    def _available_chain(self, module: str) -> list[ProviderSpec]:
        chain: list[ProviderSpec] = []
        for spec in self._slot_chain(module):
            provider = self._get_provider(spec.provider)
            if provider is not None and provider.is_available():
                chain.append(spec)
        return chain

    def _get_provider(self, name: str) -> LLMProvider | None:
        if name in self._provider_cache:
            return self._provider_cache[name]
        if not is_llm_provider_name(name, self._providers_cfg):
            return None
        cfg = self._providers_cfg.get(name, {})
        if not isinstance(cfg, dict):
            cfg = {}
        provider = create_llm_provider(name, cfg)
        if provider is None:
            return None
        self._provider_cache[name] = provider
        return provider

    def _trace(self, trace: RouterTrace | None, module: str, msg: str) -> None:
        if not self._show_trace:
            return
        persona = MODULE_PERSONAS.get(module, module)
        line = f"[{persona}] {msg}"
        if trace is not None:
            trace.add(line)

    def _record_usage(
        self,
        module: str,
        provider: str,
        model: str,
        usage: Any,
        *,
        operation: str = "llm",
        prompt_tokens: int | None = None,
        completion_tokens: int | None = None,
    ) -> None:
        if self._usage_store is None:
            return
        pt = int(prompt_tokens or 0)
        ct = int(completion_tokens or 0)
        if usage is not None:
            pt = int(getattr(usage, "prompt_tokens", pt) or pt)
            ct = int(getattr(usage, "completion_tokens", ct) or ct)
        if pt == 0 and ct == 0:
            return
        op = token_usage_operation.get()
        if operation not in ("llm",):
            op = f"{op}:{operation}" if op != "llm" else operation
        self._usage_store.record(
            user_id=token_usage_user_id.get(),
            module=module,
            operation=op,
            provider=provider,
            model=model,
            prompt_tokens=pt,
            completion_tokens=ct,
        )
