"""RoutingMixin — provider routing, reliability, fallback, circuit breaking."""
from __future__ import annotations

import asyncio
import time
from typing import Any

from logos_adk.agent._constants import _AgentRoutingDecision
from logos_adk.context import RunContext, get_run_context
from logos_adk.errors import CircuitBreakerOpenError, ConfigError, ProviderTimeoutError
from logos_adk.providers import Provider, provider_registry
from logos_adk.reliability import CircuitBreaker, ensure_circuit_allows
from logos_adk.types import Message, Response


class RoutingMixin:
    """Mixin providing provider routing and reliability methods for Agent."""

    def _estimate_prompt_complexity(self, run_context: RunContext | None) -> tuple[str, str]:
        routing = self._config.routing
        prompt = (run_context.prompt_text if run_context is not None and run_context.prompt_text else "").strip()
        task_type = (
            str(run_context.task_type).strip().lower()
            if run_context is not None and run_context.task_type
            else ""
        )
        criticality = (
            str(run_context.criticality).strip().lower()
            if run_context is not None and run_context.criticality
            else ""
        )
        routing_hint = (
            str(run_context.routing_hint).strip().lower()
            if run_context is not None and run_context.routing_hint
            else ""
        )
        normalized = prompt.lower()
        if criticality in {"critical", "high"} or routing_hint in {"critical", "quality", "strong"}:
            return "high", "explicit_criticality"
        if criticality in {"simple", "low"} or routing_hint in {"cheap", "fast", "simple"}:
            return "low", "explicit_hint"
        if task_type and task_type in {item.lower() for item in routing.critical_task_types}:
            return "high", f"critical_task_type:{task_type}"
        if task_type and task_type in {item.lower() for item in routing.simple_task_types}:
            return "low", f"simple_task_type:{task_type}"
        if any(marker.lower() in normalized for marker in routing.critical_prompt_markers):
            return "high", "critical_prompt_marker"
        if any(marker.lower() in normalized for marker in routing.simple_prompt_markers):
            return "low", "simple_prompt_marker"
        if len(normalized) >= routing.high_length_threshold or normalized.count("\n") >= 5 or normalized.count("?") >= 2:
            return "high", "prompt_length"
        if len(normalized) >= routing.medium_length_threshold or normalized.count("\n") >= 2:
            return "medium", "prompt_structure"
        return "low", "default_simple"

    def _heuristic_routing_decision(self, candidates: list[Provider]) -> _AgentRoutingDecision | None:
        if not self._config.routing.enabled or len(candidates) <= 1:
            return None
        run_context = get_run_context()
        preference = self._routing_budget_preference(run_context)
        confidence_tier, reason = self._estimate_prompt_complexity(run_context)
        pressure = self._budget_pressure(run_context)
        budget_ratio = float(pressure["budget_ratio"])
        budget_pressure = bool(pressure["near_limit"])
        if preference == "cheap" and confidence_tier != "high":
            confidence_tier = "low"
            reason = f"{reason}+budget_preference"
        elif preference == "quality":
            confidence_tier = "high"
            reason = f"{reason}+budget_preference"

        indexed = list(enumerate(candidates))
        known = [
            (index, provider, cost)
            for index, provider in indexed
            for cost in [self._provider_model_unit_cost(provider)]
            if cost is not None
        ]
        if not known:
            return None
        unknown = [(index, provider) for index, provider in indexed if self._provider_model_unit_cost(provider) is None]
        if confidence_tier == "high":
            ordered_known = sorted(known, key=lambda item: (-item[2], item[0]))
        elif confidence_tier == "low":
            ordered_known = sorted(known, key=lambda item: (item[2], item[0]))
        else:
            ordered_known = sorted(known, key=lambda item: item[0])
        if budget_pressure and preference != "quality":
            ordered_known = sorted(known, key=lambda item: (item[2], item[0]))
            runtime_status = self._runtime_budget_status(run_context)
            runtime_exhausted_scopes = [
                scope for scope, item in runtime_status.items() if item.exhausted
            ]
            runtime_near_limit_scopes = [
                scope for scope, item in runtime_status.items() if item.near_limit
            ]
            if runtime_exhausted_scopes:
                reason = f"{reason}+{'_'.join(runtime_exhausted_scopes)}_budget_exhausted"
            elif runtime_near_limit_scopes:
                reason = f"{reason}+{'_'.join(runtime_near_limit_scopes)}_budget_pressure"
            else:
                reason = f"{reason}+spend_ledger_budget_pressure"
        ordered = [provider for _, provider, _ in ordered_known] + [provider for _, provider in unknown]
        ordered_keys = [self._provider_key(provider) for provider in ordered]
        return _AgentRoutingDecision(
            ordered_model_keys=ordered_keys,
            selected_model_key=ordered_keys[0],
            confidence_tier=confidence_tier,
            budget_ratio=budget_ratio,
            routed_for_budget=budget_pressure and preference != "quality",
            source="heuristic",
            budget_preference=preference,
            reason=reason,
        )

    def _routing_decision_for_candidates(self, candidates: list[Provider]) -> Any | None:
        run_context = get_run_context()
        if self._model_routing_learning is not None and candidates:
            try:
                return self._model_routing_learning.route_candidates(
                    [self._provider_key(provider) for provider in candidates],
                    prompt=run_context.prompt_text if run_context else None,
                    assignment_key=(
                        run_context.session_id
                        or run_context.trace_id
                        if run_context is not None
                        else None
                    ),
                    workspace_id=run_context.workspace_id if run_context else None,
                    task_type=run_context.task_type if run_context else None,
                    user_segment=run_context.user_segment if run_context else None,
                )
            except Exception:
                return self._heuristic_routing_decision(candidates)
        return self._heuristic_routing_decision(candidates)

    def _ordered_candidates_from_decision(
        self,
        candidates: list[Provider],
        decision: Any | None,
    ) -> list[Provider]:
        if decision is None:
            return candidates
        by_key = {self._provider_key(provider): provider for provider in candidates}
        ordered = [by_key[key] for key in decision.ordered_model_keys if key in by_key]
        return ordered or candidates

    def _provider_candidates(self) -> list[Provider]:
        primary = self._resolve_provider()
        candidates = [primary]
        for item in self._fallback_models:
            if isinstance(item, Provider):
                candidates.append(item)
            elif isinstance(item, str):
                candidates.append(provider_registry.resolve(item))
            else:
                raise ConfigError(f"Invalid fallback model type: {type(item)}")
        routing_decision = self._routing_decision_for_candidates(candidates)
        if routing_decision is not None:
            return self._ordered_candidates_from_decision(candidates, routing_decision)
        if self._tool_workflow_learning is None or len(candidates) <= 2:
            return candidates
        run_context = get_run_context()
        fallback_providers = candidates[1:]
        ordered_keys = self._tool_workflow_learning.recommended_fallback_keys(
            [self._provider_key(provider) for provider in fallback_providers],
            task_type=run_context.task_type if run_context else None,
            workspace_id=run_context.workspace_id if run_context else None,
        )
        by_key = {self._provider_key(provider): provider for provider in fallback_providers}
        return [primary] + [by_key[key] for key in ordered_keys if key in by_key]

    def _provider_key(self, provider: Provider) -> str:
        return f"{type(provider).__name__}:{self._provider_model_name(provider)}"

    def _provider_circuit(self, provider: Provider) -> CircuitBreaker:
        key = self._provider_key(provider)
        if key not in self._provider_circuits:
            self._provider_circuits[key] = CircuitBreaker(
                failure_threshold=self._circuit_breaker_failure_threshold,
                recovery_timeout=self._circuit_breaker_recovery_timeout,
            )
        return self._provider_circuits[key]

    async def _emit_model_routing_decision_metrics(
        self,
        *,
        routing_decision: Any,
        run_context: RunContext | None,
    ) -> None:
        if routing_decision is None:
            return
        attributes = {
            "agent_name": self.name,
            "workspace_id": run_context.workspace_id if run_context else "__global__",
            "task_type": run_context.task_type if run_context else "",
            "confidence_tier": routing_decision.confidence_tier,
            "selected_model_key": routing_decision.selected_model_key,
            "used_canary": str(bool(routing_decision.used_canary)).lower(),
            "routed_for_budget": str(bool(routing_decision.routed_for_budget)).lower(),
        }
        await self._emit_metric(
            name="adk_model_routing_budget_pressure",
            value=float(routing_decision.budget_ratio),
            unit="ratio",
            attributes=attributes,
        )
        if routing_decision.used_canary:
            await self._emit_metric(
                name="adk_model_routing_canary_selected_total",
                value=1.0,
                unit="count",
                attributes=attributes,
            )

    def _routing_metadata(self, *, provider: Provider, routing_decision: Any | None) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "selected_model_key": self._provider_key(provider),
            "selected_model_name": self._provider_model_name(provider),
        }
        if routing_decision is not None:
            payload.update(
                {
                    "ordered_model_keys": list(routing_decision.ordered_model_keys),
                    "confidence_tier": routing_decision.confidence_tier,
                    "budget_ratio": float(getattr(routing_decision, "budget_ratio", 0.0) or 0.0),
                    "used_canary": bool(getattr(routing_decision, "used_canary", False)),
                    "routed_for_budget": bool(getattr(routing_decision, "routed_for_budget", False)),
                }
            )
            source = getattr(routing_decision, "source", None)
            if source:
                payload["source"] = source
            budget_preference = getattr(routing_decision, "budget_preference", None)
            if budget_preference:
                payload["budget_preference"] = budget_preference
            reason = getattr(routing_decision, "reason", None)
            if reason:
                payload["reason"] = reason
        else:
            payload["ordered_model_keys"] = [self._provider_key(provider)]
        return payload

    def _attach_routing_metadata(
        self,
        response: Response,
        *,
        provider: Provider,
        routing_decision: Any | None,
    ) -> Response:
        if routing_decision is None and not self._config.routing.enabled:
            return response
        raw = dict(response.raw or {})
        raw["routing"] = self._routing_metadata(provider=provider, routing_decision=routing_decision)
        return Response(
            content=response.content,
            content_parts=response.content_parts,
            tool_calls=response.tool_calls,
            usage=response.usage,
            raw=raw,
            reasoning=response.reasoning,
            status=response.status,
        )

    async def _emit_model_routing_result_metric(
        self,
        *,
        model_key: str,
        success: bool,
        confidence_tier: str,
        run_context: RunContext | None,
        was_canary: bool,
    ) -> None:
        await self._emit_metric(
            name="adk_model_routing_win_rate",
            value=1.0 if success else 0.0,
            unit="ratio",
            attributes={
                "agent_name": self.name,
                "workspace_id": run_context.workspace_id if run_context else "__global__",
                "task_type": run_context.task_type if run_context else "",
                "user_segment": run_context.user_segment if run_context else "",
                "model_key": model_key,
                "confidence_tier": confidence_tier,
                "was_canary": str(bool(was_canary)).lower(),
            },
        )

    async def _provider_generate_with_reliability(
        self,
        provider_messages: list[Message],
        *,
        tools: list[dict] | None,
        output_schema: Any,
        trace_id: str,
        session_id: str | None,
        request_id: str | None,
        **kwargs: Any,
    ) -> tuple[Provider, Response]:
        last_error: Exception | None = None
        primary = self._resolve_provider()
        run_context = get_run_context()
        candidates = self._provider_candidates()
        routing_decision = self._routing_decision_for_candidates(candidates)
        ordered_candidates = self._ordered_candidates_from_decision(candidates, routing_decision)
        await self._emit_model_routing_decision_metrics(
            routing_decision=routing_decision,
            run_context=run_context,
        )
        for provider in ordered_candidates:
            circuit = self._provider_circuit(provider)
            started = time.monotonic()
            provider_model_name = self._provider_model_name(provider)
            try:
                ensure_circuit_allows(circuit, type(provider).__name__)
                async def execute() -> Response:
                    try:
                        return await asyncio.wait_for(
                            provider.generate(
                                provider_messages,
                                tools=tools,
                                output_schema=output_schema,
                                **kwargs,
                            ),
                            timeout=self._provider_timeout,
                        )
                    except asyncio.TimeoutError as exc:
                        raise ProviderTimeoutError(
                            f"Provider timeout after {self._provider_timeout}s"
                        ) from exc

                response = await self._retry_policy.call(
                    execute
                )
                duration_ms = (time.monotonic() - started) * 1000
                circuit.record_success()
                self._record_spend(
                    provider=provider,
                    usage=response.usage,
                    run_context=run_context,
                    trace_id=trace_id,
                )
                if self._model_routing_learning is not None:
                    self._model_routing_learning.record_route(
                        model_key=self._provider_key(provider),
                        success=True,
                        latency_ms=duration_ms,
                        cost_usd=self._estimate_response_cost_usd(provider_model_name, response.usage),
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        workspace_id=run_context.workspace_id if run_context else None,
                        task_type=run_context.task_type if run_context else None,
                        user_segment=run_context.user_segment if run_context else None,
                        trace_id=trace_id,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                    await self._emit_model_routing_result_metric(
                        model_key=self._provider_key(provider),
                        success=True,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        run_context=run_context,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                if self._tool_workflow_learning is not None:
                    self._tool_workflow_learning.record_fallback(
                        provider_key=self._provider_key(provider),
                        success=True,
                        duration_ms=duration_ms,
                        task_type=run_context.task_type if run_context else None,
                        workspace_id=run_context.workspace_id if run_context else None,
                        trace_id=trace_id,
                        was_primary=provider is primary,
                    )
                if provider is not primary:
                    self._audit(
                        event="provider_fallback",
                        outcome="used",
                        session_id=session_id,
                        trace_id=trace_id,
                        request_id=request_id,
                        attributes={"provider": type(provider).__name__},
                    )
                return provider, self._attach_routing_metadata(
                    response,
                    provider=provider,
                    routing_decision=routing_decision,
                )
            except ProviderTimeoutError as exc:
                duration_ms = (time.monotonic() - started) * 1000
                circuit.record_failure()
                last_error = exc
                if self._model_routing_learning is not None:
                    self._model_routing_learning.record_route(
                        model_key=self._provider_key(provider),
                        success=False,
                        latency_ms=duration_ms,
                        cost_usd=0.0,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        workspace_id=run_context.workspace_id if run_context else None,
                        task_type=run_context.task_type if run_context else None,
                        user_segment=run_context.user_segment if run_context else None,
                        trace_id=trace_id,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                    await self._emit_model_routing_result_metric(
                        model_key=self._provider_key(provider),
                        success=False,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        run_context=run_context,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                if self._tool_workflow_learning is not None:
                    self._tool_workflow_learning.record_fallback(
                        provider_key=self._provider_key(provider),
                        success=False,
                        duration_ms=duration_ms,
                        task_type=run_context.task_type if run_context else None,
                        workspace_id=run_context.workspace_id if run_context else None,
                        trace_id=trace_id,
                        was_primary=provider is primary,
                        error=str(exc),
                    )
                self._audit(
                    event="provider_timeout",
                    outcome="failure",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                    attributes={"provider": type(provider).__name__},
                )
            except CircuitBreakerOpenError as exc:
                duration_ms = (time.monotonic() - started) * 1000
                last_error = exc
                if self._model_routing_learning is not None:
                    self._model_routing_learning.record_route(
                        model_key=self._provider_key(provider),
                        success=False,
                        latency_ms=duration_ms,
                        cost_usd=0.0,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        workspace_id=run_context.workspace_id if run_context else None,
                        task_type=run_context.task_type if run_context else None,
                        user_segment=run_context.user_segment if run_context else None,
                        trace_id=trace_id,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                    await self._emit_model_routing_result_metric(
                        model_key=self._provider_key(provider),
                        success=False,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        run_context=run_context,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                if self._tool_workflow_learning is not None:
                    self._tool_workflow_learning.record_fallback(
                        provider_key=self._provider_key(provider),
                        success=False,
                        duration_ms=duration_ms,
                        task_type=run_context.task_type if run_context else None,
                        workspace_id=run_context.workspace_id if run_context else None,
                        trace_id=trace_id,
                        was_primary=provider is primary,
                        error=str(exc),
                    )
                continue
            except Exception as exc:
                duration_ms = (time.monotonic() - started) * 1000
                circuit.record_failure()
                last_error = exc
                if self._model_routing_learning is not None:
                    self._model_routing_learning.record_route(
                        model_key=self._provider_key(provider),
                        success=False,
                        latency_ms=duration_ms,
                        cost_usd=0.0,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        workspace_id=run_context.workspace_id if run_context else None,
                        task_type=run_context.task_type if run_context else None,
                        user_segment=run_context.user_segment if run_context else None,
                        trace_id=trace_id,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                    await self._emit_model_routing_result_metric(
                        model_key=self._provider_key(provider),
                        success=False,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        run_context=run_context,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                if self._tool_workflow_learning is not None:
                    self._tool_workflow_learning.record_fallback(
                        provider_key=self._provider_key(provider),
                        success=False,
                        duration_ms=duration_ms,
                        task_type=run_context.task_type if run_context else None,
                        workspace_id=run_context.workspace_id if run_context else None,
                        trace_id=trace_id,
                        was_primary=provider is primary,
                        error=str(exc),
                    )
        if last_error is None:
            raise RuntimeError("No provider candidates available")
        raise last_error
