# -*- coding: utf-8 -*-
"""Provider-neutral model routing primitives for SmartPaper.

Legacy tasks continue to route by coarse capability.  A5-R3-R4B adds a
separate, trusted product-policy seam: academic callers provide an explicit
logical_node, while the frozen TaskPolicySnapshot supplies provider/model and
execution settings.  Product routing never falls back to legacy capability
routes.
"""
from __future__ import annotations

from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Callable, Iterable


class ModelGatewayError(RuntimeError):
    pass


class ModelProviderUnavailableError(ModelGatewayError):
    pass


class ModelAuthenticationError(ModelGatewayError):
    pass


class ModelBalanceError(ModelGatewayError):
    pass


class ModelRequestError(ModelGatewayError):
    pass


@dataclass(frozen=True)
class ModelRoute:
    provider: str
    provider_role: str
    fallback_providers: tuple[str, ...] = ()


@dataclass(frozen=True)
class TrustedProductRoute:
    provider: str
    model_id: str
    thinking_enabled: bool
    reasoning_effort: str | None
    max_output_tokens: int
    retry_budget: int


class ProviderNeutralModelGateway:
    """Route legacy capabilities or trusted product logical nodes to adapters."""

    def __init__(self, config: dict[str, Any] | None = None):
        cfg = dict(config or {})
        self.default_provider = str(cfg.get("default_provider") or "deepseek").strip().lower()
        raw_routes = dict(cfg.get("capability_routes") or {})
        self._routes: dict[str, ModelRoute] = {}
        for capability, row in raw_routes.items():
            cap = str(capability or "").strip().lower()
            if not cap:
                continue
            if isinstance(row, str):
                provider = str(row or self.default_provider).strip().lower()
                provider_role = cap
                fallbacks = ()
            else:
                item = dict(row or {})
                provider = str(item.get("provider") or self.default_provider).strip().lower()
                provider_role = str(item.get("provider_role") or cap).strip().lower()
                raw_fallbacks = item.get("fallback_providers") or ()
                if isinstance(raw_fallbacks, str):
                    raw_fallbacks = [part.strip() for part in raw_fallbacks.split(",") if part.strip()]
                fallbacks = tuple(
                    str(name or "").strip().lower()
                    for name in raw_fallbacks
                    if str(name or "").strip()
                )
            self._routes[cap] = ModelRoute(
                provider=provider,
                provider_role=provider_role,
                fallback_providers=fallbacks,
            )
        self._call_adapters: dict[str, Callable[..., str]] = {}
        self._stream_adapters: dict[str, Callable[..., Iterable[str]]] = {}
        self._policy_snapshot: dict[str, Any] = {}
        self._trusted_product_policy_snapshot: dict[str, Any] = {}

    def register_provider(
        self,
        provider: str,
        *,
        call_adapter: Callable[..., str],
        stream_adapter: Callable[..., Iterable[str]],
    ) -> None:
        name = str(provider or "").strip().lower()
        if not name:
            raise ValueError("provider name is required")
        self._call_adapters[name] = call_adapter
        self._stream_adapters[name] = stream_adapter

    # Legacy config/request policy.  It remains a separate trust domain.
    def set_policy_snapshot(self, snapshot: dict[str, Any] | None) -> None:
        self._policy_snapshot = deepcopy(dict(snapshot or {}))

    @property
    def policy_snapshot(self) -> dict[str, Any]:
        return deepcopy(self._policy_snapshot)

    def set_trusted_product_policy_snapshot(self, snapshot: dict[str, Any] | None) -> None:
        """Attach a Worker/D1-verified TaskPolicySnapshot for product execution.

        The runner validates task identity and snapshot SHA before calling this
        setter.  The gateway still validates route shape and fails closed before
        invoking any provider adapter.
        """
        self._trusted_product_policy_snapshot = deepcopy(dict(snapshot or {}))

    @property
    def trusted_product_policy_snapshot(self) -> dict[str, Any]:
        return deepcopy(self._trusted_product_policy_snapshot)

    def _route(self, capability: str) -> ModelRoute:
        cap = str(capability or "reasoning").strip().lower() or "reasoning"
        policy_routes = dict(self._policy_snapshot.get("model_routes") or {})
        row = policy_routes.get(cap)
        if isinstance(row, str):
            return ModelRoute(provider=row.strip().lower(), provider_role=cap)
        if isinstance(row, dict):
            raw_fallbacks = row.get("fallback_providers") or ()
            if isinstance(raw_fallbacks, str):
                raw_fallbacks = [part.strip() for part in raw_fallbacks.split(",") if part.strip()]
            return ModelRoute(
                provider=str(row.get("provider") or self.default_provider).strip().lower(),
                provider_role=str(row.get("provider_role") or cap).strip().lower(),
                fallback_providers=tuple(
                    str(name or "").strip().lower()
                    for name in raw_fallbacks
                    if str(name or "").strip()
                ),
            )
        return self._routes.get(cap) or ModelRoute(self.default_provider, cap)

    @staticmethod
    def _provider_sequence(route: ModelRoute) -> tuple[str, ...]:
        return tuple(dict.fromkeys(
            name for name in (route.provider, *route.fallback_providers) if name
        ))

    def _trusted_product_context(self) -> tuple[dict[str, Any], dict[str, Any]] | None:
        snapshot = self._trusted_product_policy_snapshot
        if not snapshot:
            return None
        resolved = snapshot.get("resolved_policy")
        if not isinstance(resolved, dict):
            raise ModelRequestError("trusted_product_policy_invalid")
        model = resolved.get("model")
        if not isinstance(model, dict):
            raise ModelRequestError("trusted_product_model_policy_missing")
        routes = model.get("routes")
        if not isinstance(routes, dict) or not routes:
            raise ModelRequestError("trusted_product_model_routes_missing")
        return resolved, routes

    def _trusted_product_route(self, logical_node: str | None) -> tuple[TrustedProductRoute, str]:
        context = self._trusted_product_context()
        if context is None:
            raise ModelRequestError("trusted_product_policy_missing")
        resolved, routes = context
        node = str(logical_node or "").strip()
        if not node:
            raise ModelRequestError("trusted_product_logical_node_required")
        row = routes.get(node)
        if not isinstance(row, dict):
            raise ModelRequestError(f"trusted_product_route_missing:{node}")

        provider = str(row.get("provider") or "").strip().lower()
        model_id = str(row.get("model_id_or_alias") or "").strip()
        if not provider or not model_id:
            raise ModelRequestError(f"trusted_product_route_invalid:{node}")

        thinking = row.get("thinking_enabled")
        if not isinstance(thinking, bool):
            raise ModelRequestError(f"trusted_product_route_thinking_invalid:{node}")
        effort = row.get("reasoning_effort")
        if thinking:
            if str(effort or "").strip().lower() != "high":
                raise ModelRequestError(f"trusted_product_route_effort_invalid:{node}")
            effort = "high"
        else:
            if effort not in (None, ""):
                raise ModelRequestError(f"trusted_product_route_effort_invalid:{node}")
            effort = None

        try:
            max_output_tokens = int(row.get("max_output_tokens"))
            retry_budget = int(row.get("retry_budget"))
        except (TypeError, ValueError):
            raise ModelRequestError(f"trusted_product_route_budget_invalid:{node}")
        if max_output_tokens <= 0 or retry_budget < 0:
            raise ModelRequestError(f"trusted_product_route_budget_invalid:{node}")
        if row.get("escalation_model_id_or_alias") not in (None, ""):
            raise ModelRequestError(f"trusted_product_route_escalation_forbidden:{node}")
        if row.get("escalation_condition_code") not in (None, ""):
            raise ModelRequestError(f"trusted_product_route_escalation_forbidden:{node}")

        model_policy_version = str(
            resolved.get("model_policy_version")
            or snapshot_model_policy_version(self._trusted_product_policy_snapshot)
            or ""
        ).strip()
        return TrustedProductRoute(
            provider=provider,
            model_id=model_id,
            thinking_enabled=thinking,
            reasoning_effort=effort,
            max_output_tokens=max_output_tokens,
            retry_budget=retry_budget,
        ), model_policy_version

    @staticmethod
    def _effective_max_tokens(requested: Any, ceiling: int) -> int:
        if requested is None:
            return int(ceiling)
        try:
            value = int(requested)
        except (TypeError, ValueError):
            raise ModelRequestError("trusted_product_max_tokens_invalid")
        if value <= 0:
            raise ModelRequestError("trusted_product_max_tokens_invalid")
        return min(value, int(ceiling))

    def _product_payload(self, *, capability: str, logical_node: str | None, kwargs: dict[str, Any]):
        route, model_policy_version = self._trusted_product_route(logical_node)
        payload = dict(kwargs)
        payload["model_role"] = str(capability or "reasoning").strip().lower() or "reasoning"
        payload["model_name"] = route.model_id
        payload["thinking"] = route.thinking_enabled
        payload["reasoning_effort"] = route.reasoning_effort
        payload["allow_fallback"] = False
        payload["max_tokens"] = self._effective_max_tokens(
            payload.get("max_tokens"), route.max_output_tokens
        )
        payload["retry_budget"] = route.retry_budget
        payload["logical_node"] = str(logical_node or "").strip()
        payload["model_policy_version"] = model_policy_version
        payload["route_enforced"] = True
        return route.provider, payload

    def call(self, *, capability: str, logical_node: str | None = None, **kwargs) -> str:
        if self._trusted_product_context() is not None:
            provider, payload = self._product_payload(
                capability=capability, logical_node=logical_node, kwargs=kwargs
            )
            adapter = self._call_adapters.get(provider)
            if not callable(adapter):
                raise ModelProviderUnavailableError(
                    f"model_provider_unavailable:{provider}:{logical_node or capability}"
                )
            return adapter(**payload)

        route = self._route(capability)
        missing = []
        for provider in self._provider_sequence(route):
            adapter = self._call_adapters.get(provider)
            if not callable(adapter):
                missing.append(provider)
                continue
            payload = dict(kwargs)
            payload["model_role"] = route.provider_role
            return adapter(**payload)
        raise ModelProviderUnavailableError(
            f"model_provider_unavailable:{','.join(missing) or route.provider}:{capability}"
        )

    def stream(self, *, capability: str, logical_node: str | None = None, **kwargs):
        if self._trusted_product_context() is not None:
            provider, payload = self._product_payload(
                capability=capability, logical_node=logical_node, kwargs=kwargs
            )
            adapter = self._stream_adapters.get(provider)
            if not callable(adapter):
                raise ModelProviderUnavailableError(
                    f"model_provider_unavailable:{provider}:{logical_node or capability}"
                )
            yield from adapter(**payload)
            return

        route = self._route(capability)
        missing = []
        for provider in self._provider_sequence(route):
            adapter = self._stream_adapters.get(provider)
            if not callable(adapter):
                missing.append(provider)
                continue
            payload = dict(kwargs)
            payload["model_role"] = route.provider_role
            yield from adapter(**payload)
            return
        raise ModelProviderUnavailableError(
            f"model_provider_unavailable:{','.join(missing) or route.provider}:{capability}"
        )


def snapshot_model_policy_version(snapshot: dict[str, Any] | None) -> str:
    data = dict(snapshot or {})
    return str(data.get("model_policy_version") or "").strip()


def resolve_plan_policy_snapshot(
    app_config: dict[str, Any], request: dict[str, Any] | None = None
) -> dict[str, Any]:
    """Freeze a config-driven legacy plan policy without touching payment authority."""
    req = dict(request or {})
    policies = dict(app_config.get("plan_policies") or {})
    entitlement = req.get("entitlement_snapshot") if isinstance(req.get("entitlement_snapshot"), dict) else {}
    requested = str(
        req.get("trusted_plan_id") or entitlement.get("plan_id") or ""
    ).strip().lower()
    selected = requested if requested in policies else "legacy_current"
    policy = deepcopy(dict(policies.get(selected) or policies.get("legacy_current") or {}))
    policy["plan_id"] = selected
    policy["policy_version"] = str(
        policy.get("policy_version") or app_config.get("plan_policy_version") or "plan-policy-v1"
    )
    return policy