"""Model Gateway — unified interface for the Model Layer.

The Model Gateway provides a single entry point for:
- Chat/completion requests with automatic routing and fallback
- Embedding generation
- Model and provider registry management
- Model intelligence (discovery, validation, classification)
- Health checking across all providers
"""

from __future__ import annotations

import hashlib
import json
import logging
from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import Any

from cyberbrain26.control_plane.model_gateway.model_egress import ModelEgressGate
from cyberbrain26.model_layer import (
    LiteLLMClient,
    LiteLLMConfig,
    ModelIntelligence,
    ModelRecord,
    ModelRegistry,
    ModelRouter,
    OllamaClient,
    OllamaConfig,
    ProviderRecord,
    ProviderRegistry,
    ProviderStatus,
    ProviderTrust,
    RouterConfig,
)
from cyberbrain26.model_layer.llama_cpp_client import LlamaCppClient, LlamaCppConfig
from cyberbrain26.model_layer.provider_registry import ProviderStatus as _ProviderStatus
from cyberbrain26.security_plane.permission import SecurityError

logger = logging.getLogger(__name__)


@dataclass
class ModelGatewayConfig:
    """Configuration for Model Gateway."""

    default_provider: str = "litellm"
    fallback_providers: list[str] = field(default_factory=lambda: ["ollama", "llama.cpp"])
    prefer_local: bool = True
    prefer_free: bool = True
    free_only: bool = True
    strategy: str = "local_first"

    litellm: LiteLLMConfig | None = None
    ollama: OllamaConfig | None = None
    llama_cpp: LlamaCppConfig | None = None
    router: RouterConfig | None = None


class ModelGatewayError(Exception):
    """Base exception for Model Gateway errors."""


class ModelNotAvailableError(ModelGatewayError):
    """Raised when no suitable model is available."""


class ModelGateway:
    """Unified Model Gateway for CyberBrain26."""

    def __init__(
        self,
        config: ModelGatewayConfig | None = None,
        *,
        egress_gate: Any | None = None,
        agent_id: str = "model-gateway",
    ) -> None:
        self.config = config or ModelGatewayConfig()

        self.litellm_client = LiteLLMClient(self.config.litellm)
        self.ollama_client = OllamaClient(self.config.ollama)
        self.llama_cpp_client = LlamaCppClient(self.config.llama_cpp)

        self.model_registry: ModelRegistry = ModelRegistry()
        self.provider_registry: ProviderRegistry = ProviderRegistry()
        self.model_intelligence = ModelIntelligence(
            model_registry=self.model_registry,
            provider_registry=self.provider_registry,
        )

        self.router = ModelRouter(
            model_registry=self.model_registry,
            provider_registry=self.provider_registry,
            config=self.config.router or RouterConfig(
                prefer_local=self.config.prefer_local,
                prefer_free=self.config.prefer_free,
            ),
            model_intelligence=self.model_intelligence,
        )

        self._register_builtin_providers()

        # When an EgressGate is injected, every chat() call must clear it
        # before a payload reaches a provider.  Without a gate the gateway
        # keeps its previous behaviour, so wiring is opt-in and the security
        # boundary is explicit rather than implied.
        self.agent_id = agent_id
        self.model_egress = ModelEgressGate(egress_gate=egress_gate)

    def _register_builtin_providers(self) -> None:
        litellm_provider = ProviderRecord(
            provider="litellm",
            status=ProviderStatus.ACTIVE,
            free_tier=True,
            registration_required=False,
            card_required=False,
            api_base="https://api.litellm.ai",
            api_key_required=True,
            api_available=True,
            available_models=[
                "gemini-flash",
                "gemini-pro",
                "gpt-4o",
                "gpt-4o-mini",
                "claude-3-5-sonnet",
                "claude-3-haiku",
                "mistral-large",
                "mistral-medium",
                "groq-llama-3",
                "groq-mixtral",
            ],
            trust=ProviderTrust.OFFICIAL,
            source="litellm",
            free=True,
            no_card=True,
            api_availability=True,
            official_source=True,
        )
        self.provider_registry.register(litellm_provider)

        ollama_provider = ProviderRecord(
            provider="ollama",
            status=ProviderStatus.ACTIVE,
            free_tier=True,
            registration_required=False,
            card_required=False,
            api_base="http://127.0.0.1:11434",
            api_key_required=False,
            api_available=True,
            available_models=[
                "llama3",
                "llama3.1",
                "mistral",
                "mixtral",
                "phi3",
                "gemma2",
                "qwen2.5",
            ],
            trust=ProviderTrust.COMMUNITY,
            source="ollama",
            free=True,
            no_card=True,
            local_availability=True,
            api_availability=True,
        )
        self.provider_registry.register(ollama_provider)

        llama_cpp_provider = ProviderRecord(
            provider="llama.cpp",
            status=ProviderStatus.ACTIVE,
            free_tier=True,
            registration_required=False,
            card_required=False,
            api_base="http://127.0.0.1:8080",
            api_key_required=False,
            api_available=True,
            available_models=["llama.cpp"],
            trust=ProviderTrust.COMMUNITY,
            source="llama.cpp",
            free=True,
            no_card=True,
            local_availability=True,
            api_availability=True,
        )
        self.provider_registry.register(llama_cpp_provider)

        self._register_builtin_models()

    def _register_builtin_models(self) -> None:
        from cyberbrain26.model_layer import (
            ModelCapabilities,
            ModelFreeStatus,
            ModelRecord,
            ModelStatus,
            ModelTrust,
            ModelType,
        )

        for model_id in (self.provider_registry.get("litellm") or ProviderRecord(provider="litellm", available_models=[])).available_models:
            self.model_registry.register(
                ModelRecord(
                    model_id=model_id,
                    provider="litellm",
                    model_type=ModelType.CHAT,
                    status=ModelStatus.ACTIVE,
                    api_available=True,
                    local_available=False,
                    free_status=ModelFreeStatus.FREE,
                    trust=ModelTrust.OFFICIAL,
                    capabilities=ModelCapabilities(
                        tool_calling=True,
                        reasoning=True,
                        coding=True,
                    ),
                )
            )

        for model_id in (self.provider_registry.get("ollama") or ProviderRecord(provider="ollama", available_models=[])).available_models:
            self.model_registry.register(
                ModelRecord(
                    model_id=model_id,
                    provider="ollama",
                    model_type=ModelType.CHAT,
                    status=ModelStatus.ACTIVE,
                    api_available=True,
                    local_available=True,
                    free_status=ModelFreeStatus.FREE,
                    trust=ModelTrust.COMMUNITY,
                    capabilities=ModelCapabilities(
                        tool_calling=True,
                        coding=True,
                    ),
                )
            )

        for model_id in (self.provider_registry.get("llama.cpp") or ProviderRecord(provider="llama.cpp", available_models=[])).available_models:
            self.model_registry.register(
                ModelRecord(
                    model_id=model_id,
                    provider="llama.cpp",
                    model_type=ModelType.CHAT,
                    status=ModelStatus.ACTIVE,
                    api_available=True,
                    local_available=True,
                    free_status=ModelFreeStatus.FREE,
                    trust=ModelTrust.COMMUNITY,
                    capabilities=ModelCapabilities(),
                )
            )

    def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        stream: bool = False,
        prefer_local: bool | None = None,
        fallback: bool = True,
        classifications: frozenset[Any] | None = None,
    ) -> str | Iterator[dict[str, Any]]:
        route = self._resolve_route(model_id=model, prefer_local=prefer_local)

        self._authorize_egress(route.provider, model, messages, classifications)

        try:
            return self._execute_chat(
                route, model, messages, temperature, max_tokens, stream,
            )
        except Exception as exc:
            if fallback:
                logger.warning("Primary model failed, trying fallback: %s", exc)
                return self._fallback_chat(
                    model, messages, temperature, max_tokens, stream,
                    classifications=classifications,
                )
            raise ModelNotAvailableError(
                f"Chat failed for model {model}: {exc}"
            ) from exc

    def _model_is_free(self, model_id: str) -> bool:
        """Whether the registry records ``model_id`` as free-eligible.

        Unknown models are treated as not free: the free-only guarantee must
        fail closed when the registry has no record to consult.
        """
        record = self.model_registry.get(model_id)
        return bool(record is not None and record.is_free)

    def _authorize_egress(
        self,
        provider_name: str,
        model: str,
        messages: list[dict[str, str]],
        classifications: frozenset[Any] | None,
    ) -> None:
        """Clear the request with the egress gate before any provider call.

        Fail-closed: a missing classification is refused rather than assumed
        public.  No gate configured means no check, which is the pre-F4
        behaviour and is explicit at construction time.
        """
        if self.model_egress._egress_gate is None:
            return

        if not classifications:
            raise SecurityError(
                "model egress requires an explicit classification; refusing to "
                "assume PUBLIC"
            )

        payload = json.dumps(messages, ensure_ascii=False, sort_keys=True)
        self.model_egress.authorize_model_call(
            agent_id=self.agent_id,
            provider_id=provider_name,
            model_id=model,
            payload_sha256=hashlib.sha256(payload.encode("utf-8")).hexdigest(),
            classifications=classifications,
        )

    def _resolve_route(
        self,
        *,
        model_id: str | None = None,
        prefer_local: bool | None = None,
    ) -> ProviderRecord:
        if model_id:
            model = self.model_registry.get(model_id)
            if model and model.local_available:
                provider = self.provider_registry.get(model.provider)
                if provider:
                    return provider

        if prefer_local or self.config.prefer_local:
            local_providers = self.provider_registry.list_local()
            for p in local_providers:
                if p.status == _ProviderStatus.ACTIVE:
                    return p

        litellm = self.provider_registry.get("litellm")
        if litellm and litellm.status == _ProviderStatus.ACTIVE:
            return litellm

        active = self.provider_registry.list_active()
        if active:
            return active[0]

        raise ModelNotAvailableError("No active providers available")

    def _execute_chat(
        self,
        provider: ProviderRecord,
        model: str,
        messages: list[dict[str, str]],
        temperature: float | None,
        max_tokens: int | None,
        stream: bool,
    ) -> str | Iterator[dict[str, Any]]:
        if provider.provider == "litellm":
            return self.litellm_client.chat(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=stream,
            )
        elif provider.provider == "ollama":
            return self.ollama_client.chat(
                model=model,
                messages=messages,
                temperature=temperature,
                num_predict=max_tokens,
                stream=stream,
            )
        elif provider.provider == "llama.cpp":
            return self.llama_cpp_client.chat(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=stream,
            )
        else:
            raise ModelNotAvailableError(f"Unknown provider: {provider.provider}")

    def _fallback_chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        temperature: float | None,
        max_tokens: int | None,
        stream: bool,
        classifications: frozenset[Any] | None = None,
    ) -> str | Iterator[dict[str, Any]]:
        fallback_providers = self.config.fallback_providers
        for provider_name in fallback_providers:
            provider = self.provider_registry.get(provider_name)
            if provider and provider.status == _ProviderStatus.ACTIVE:
                # A fallback switches provider, not model, so the free-only
                # guarantee has to be re-checked here: this loop never touches
                # ModelRouter, which is where free_only is otherwise enforced.
                if self.config.free_only and not self._model_is_free(model):
                    logger.warning(
                        "Fallback provider %s refused: model %r is not free-eligible",
                        provider_name,
                        model,
                    )
                    continue
                try:
                    self._authorize_egress(
                        provider.provider, model, messages, classifications
                    )
                except SecurityError:
                    raise
                except Exception:  # noqa: BLE001
                    logger.warning(
                        "Fallback provider %s blocked by egress gate", provider_name
                    )
                    continue
                try:
                    return self._execute_chat(
                        provider, model, messages, temperature, max_tokens, stream,
                    )
                except Exception:  # noqa: BLE001
                    logger.warning(
                        "Fallback provider %s failed", provider_name
                    )
                    continue

        raise ModelNotAvailableError(
            "All fallback providers exhausted for model: " + model
        )

    def embed(
        self,
        model: str,
        texts: list[str],
        *,
        prefer_local: bool | None = None,
    ) -> list[list[float]]:
        if prefer_local or self.config.prefer_local:
            local_providers = self.provider_registry.list_local()
            for p in local_providers:
                if p.status == _ProviderStatus.ACTIVE:
                    return self._execute_embed(p, model, texts)

        litellm = self.provider_registry.get("litellm")
        if litellm and litellm.status == _ProviderStatus.ACTIVE:
            return self._execute_embed(litellm, model, texts)

        raise ModelNotAvailableError("No provider for embeddings")

    def _execute_embed(
        self,
        provider: ProviderRecord,
        model: str,
        texts: list[str],
    ) -> list[list[float]]:
        if provider.provider == "litellm":
            return self.litellm_client.embed(model=model, texts=texts)
        elif provider.provider == "ollama":
            return self.ollama_client.embed(model=model, texts=texts)
        elif provider.provider == "llama.cpp":
            return self.llama_cpp_client.embed(model=model, texts=texts)
        else:
            raise ModelNotAvailableError(f"Unknown provider: {provider.provider}")

    def register_model(self, model: ModelRecord) -> ModelRecord:
        return self.model_registry.register(model)

    def get_model(self, model_id: str) -> ModelRecord | None:
        return self.model_registry.get(model_id)

    def list_models(self) -> list[ModelRecord]:
        return self.model_registry.list_all()

    def register_provider(self, provider: ProviderRecord) -> ProviderRecord:
        return self.provider_registry.register(provider)

    def get_provider(self, provider_name: str) -> ProviderRecord | None:
        return self.provider_registry.get(provider_name)

    def list_providers(self) -> list[ProviderRecord]:
        return self.provider_registry.list_all()

    def discover_models(self) -> list[ModelRecord]:
        return self.model_intelligence.discover_from_providers()

    def run_validation_pipeline(self, models: list[ModelRecord] | None = None) -> list[dict[str, Any]]:
        return self.model_intelligence.run_pipeline(models)

    def classify_model(self, model: ModelRecord) -> str:
        return self.model_intelligence.classify(model).value

    def score_model(self, model: ModelRecord) -> float:
        return self.model_intelligence.score_model(model)

    def route(
        self,
        *,
        model_id: str | None = None,
        capabilities: dict[str, bool] | None = None,
        prefer_local: bool | None = None,
        prefer_free: bool | None = None,
        strategy: str | None = None,
        free_only: bool | None = None,
    ) -> dict[str, Any]:
        from cyberbrain26.model_layer import RouteStrategy

        strat = (
            RouteStrategy(strategy)
            if strategy
            else RouteStrategy(self.config.strategy)
        )

        route = self.router.route(
            model_id=model_id,
            capabilities=capabilities,
            prefer_local=prefer_local,
            prefer_free=prefer_free,
            strategy=strat,
            free_only=self.config.free_only if free_only is None else free_only,
        )

        return {
            "model_id": route.model_id,
            "provider": route.provider,
            "strategy": route.strategy.value,
            "score": route.score,
            "fallback": route.fallback,
            "reason": route.reason,
        }

    def route_with_fallback(
        self,
        *,
        model_id: str | None = None,
        capabilities: dict[str, bool] | None = None,
        free_only: bool | None = None,
    ) -> list[dict[str, Any]]:
        routes = self.router.route_with_fallback(
            model_id=model_id,
            capabilities=capabilities,
            free_only=self.config.free_only if free_only is None else free_only,
        )
        return [
            {
                "model_id": r.model_id,
                "provider": r.provider,
                "strategy": r.strategy.value,
                "score": r.score,
                "fallback": r.fallback,
                "reason": r.reason,
            }
            for r in routes
        ]

    def health_check(self) -> dict[str, Any]:
        results: dict[str, Any] = {
            "providers": {},
            "models": {
                "total": self.model_registry.count(),
                "active": len(self.model_registry.list_active()),
                "local": len(self.model_registry.list_local()),
                "free": len(self.model_registry.list_free()),
            },
            "routing": {
                "circuit_breakers": self.router.get_circuit_breaker_status(),
                "total_routes": len(self.router.get_route_history()),
            },
        }

        for provider in self.provider_registry.list_all():
            client = self._get_client(provider.provider)
            if client:
                results["providers"][provider.provider] = client.health_check()
            else:
                results["providers"][provider.provider] = {
                    "healthy": False,
                    "provider": provider.provider,
                    "error": "no client",
                }

        return results

    def _get_client(self, provider_name: str) -> Any:
        if provider_name == "litellm":
            return self.litellm_client
        elif provider_name == "ollama":
            return self.ollama_client
        elif provider_name == "llama.cpp":
            return self.llama_cpp_client
        return None

    def is_healthy(self) -> bool:
        for provider in self.provider_registry.list_active():
            client = self._get_client(provider.provider)
            if client and client.is_available():
                return True
        return False