"""Centralized AI gateway environment-driven foundation for Batch 2.

This module follows the approved AI gateway environment-driven concept for Batch 2:
- task-based routing policy driven by PROJECT_STATUS (development vs production);
- explicit task contracts for LLM and Embedding routing;
- separate auditable key pools for OpenAI and Dashscope;
- provider-agnostic embedding path fixed to 1536 dimensions;
- intra-provider LLM fallback only for the active profile provider;
- typed gateway errors for exhaustion, provider failure, and validation.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

try:
    from langchain_openai import ChatOpenAI
except ImportError:  # pragma: no cover - handled by runtime validation.
    ChatOpenAI = None

try:
    from openai import AsyncOpenAI, OpenAI
except ImportError:  # pragma: no cover - handled by runtime validation.
    AsyncOpenAI = None
    from openai import OpenAI

from app.core.config import get_settings

logger = logging.getLogger(__name__)

EMBEDDING_DIMENSION = 1536

ROUTING_POLICY: dict[str, dict[str, Any]] = {
    "development": {
        "llm_provider": "dashscope",
        "embedding_provider": "dashscope",
        "embeddings": {
            "default": "text-embedding-v4",
        },
        "providers": {
            "openai": {
                "base_url": None,
            },
            "dashscope": {
                "base_url": "https://ws-3zxwxknmp91r3atr.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1",
            },
        },
        "tasks": {
            "chatbot": {
                "models": ["qwen-max", "qwen-plus", "qwen-turbo"],
                "params": {"temperature": 0.3},
            },
            "ingestion": {
                "models": ["qwen-plus", "qwen-turbo"],
                "params": {"temperature": 0.0},
            },
            "comparative": {
                "models": ["qwen-turbo"],
                "params": {"temperature": 0.1},
            },
        },
    },
    "production": {
        "llm_provider": "openai",
        "embedding_provider": "openai",
        "embeddings": {
            "default": "text-embedding-3-small",
        },
        "providers": {
            "openai": {
                "base_url": None,
            },
            "dashscope": {
                "base_url": "https://ws-3zxwxknmp91r3atr.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1",
            },
        },
        "tasks": {
            "chatbot": {
                "models": ["gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano"],
                "params": {"temperature": 0.3},
            },
            "ingestion": {
                "models": ["gpt-5.4-mini", "gpt-5.4-nano"],
                "params": {"temperature": 0.0},
            },
            "comparative": {
                "models": ["gpt-5.4-nano"],
                "params": {"temperature": 0.1},
            },
        },
    },
}


def _build_legacy_profiles(
    routing_policy: dict[str, dict[str, Any]],
) -> dict[str, dict[str, Any]]:
    """Derive the legacy profile view from the task-based routing policy."""
    return {
        profile_name: {
            "llm_provider": profile_config["llm_provider"],
            "embedding_provider": profile_config["embedding_provider"],
            "embedding_model": profile_config["embeddings"]["default"],
        }
        for profile_name, profile_config in routing_policy.items()
    }


def _build_legacy_providers(
    routing_policy: dict[str, dict[str, Any]],
) -> dict[str, dict[str, Any]]:
    """Derive the legacy provider view from the task-based routing policy."""
    provider_registry: dict[str, dict[str, Any]] = {}
    for profile_config in routing_policy.values():
        tasks = profile_config.get("tasks", {})
        for provider_name, provider_config in profile_config.get("providers", {}).items():
            provider_entry = provider_registry.setdefault(
                provider_name,
                {
                    "llm_models": [],
                    "base_url": provider_config.get("base_url"),
                },
            )
            if provider_entry["base_url"] != provider_config.get("base_url"):
                raise ValueError(
                    f"Provider base_url mismatch detected for '{provider_name}'."
                )

            if provider_name != profile_config["llm_provider"]:
                continue

            for task_config in tasks.values():
                for model_name in task_config.get("models", []):
                    if model_name not in provider_entry["llm_models"]:
                        provider_entry["llm_models"].append(model_name)

    return provider_registry


PROFILES: dict[str, dict[str, Any]] = _build_legacy_profiles(ROUTING_POLICY)
PROVIDERS: dict[str, dict[str, Any]] = _build_legacy_providers(ROUTING_POLICY)


class AIGatewayError(RuntimeError):
    """Base error for gateway-controlled runtime failures."""


class AIGatewayInitializationError(AIGatewayError):
    """Raised when gateway initialization cannot satisfy target-state rules."""


class AIGatewayValidationError(AIGatewayError):
    """Raised when a caller sends an invalid gateway request."""


class ProviderRequestError(AIGatewayError):
    """Raised when an active provider request fails without a retry path."""


class EmbeddingPathExhaustedError(AIGatewayError):
    """Raised when all embedding keys are exhausted for the fixed model path."""


class LLMPathExhaustedError(AIGatewayError):
    """Raised when all keys and models fail in the active provider path."""


class LangChainPathExhaustedError(AIGatewayError):
    """Raised when all LangChain fallback candidates fail."""


@dataclass(frozen=True)
class KeyPoolState:
    """Auditable snapshot of a key pool.

    Attributes:
        pool_name: Logical pool identifier.
        source: Environment contract backing this pool.
        size: Number of keys configured in the pool.
        current_index: Current pointer inside the pool.
        exhausted: Whether all configured keys have been consumed.
    """

    pool_name: str
    source: str
    size: int
    current_index: int
    exhausted: bool


class KeyPoolManager:
    """Explicit, ordered key pool manager for target-state gateway rotation."""

    def __init__(self, pool_name: str, source: str, keys: tuple[str, ...]) -> None:
        """Initialize a key pool from normalized key values.

        Args:
            pool_name: Logical key pool name.
            source: Environment contract source.
            keys: Ordered keys for the pool.
        """
        self._pool_name = pool_name
        self._source = source
        self._keys = keys
        self._current_index = 0

    @property
    def state(self) -> KeyPoolState:
        """Return the current auditable pool state."""
        return KeyPoolState(
            pool_name=self._pool_name,
            source=self._source,
            size=len(self._keys),
            current_index=self._current_index,
            exhausted=self.is_exhausted,
        )

    @property
    def is_empty(self) -> bool:
        """Return whether the pool contains no configured keys."""
        return len(self._keys) == 0

    @property
    def is_exhausted(self) -> bool:
        """Return whether the current index has consumed every key."""
        return self.is_empty or self._current_index >= len(self._keys)

    def current_key(self) -> str:
        """Expose the current key without mutating pool state.

        Returns:
            str: Current key.

        Raises:
            AIGatewayError: If the pool has no remaining keys.
        """
        if self.is_exhausted:
            raise AIGatewayError(f"Key pool '{self._pool_name}' is exhausted.")
        return self._keys[self._current_index]

    def rotate(self) -> bool:
        """Advance to the next key in deterministic order.

        Returns:
            bool: True if a next key exists, otherwise False.
        """
        if self.is_exhausted:
            return False
        self._current_index += 1
        return not self.is_exhausted

    def reset(self) -> None:
        """Reset the pool pointer to the first configured key."""
        self._current_index = 0


class EmbeddingEngine:
    """Environment-driven embedding engine supporting multiple providers."""

    def __init__(
        self,
        embedding_pool: KeyPoolManager,
        provider_name: str,
        model_name: str,
        base_url: str | None = None,
    ) -> None:
        """Initialize the embedding engine for a specific provider.

        Args:
            embedding_pool: The key pool for the active provider.
            provider_name: The active provider name.
            model_name: The active model name.
            base_url: Optional base URL for API compatibility.

        Raises:
            AIGatewayInitializationError: If the embedding pool is empty.
        """
        if embedding_pool.is_empty:
            raise AIGatewayInitializationError(
                f"Embedding key pool for '{provider_name}' must contain at least one key."
            )
        self._embedding_pool = embedding_pool
        self._provider_name = provider_name
        self._model_name = model_name
        self._base_url = base_url

    @property
    def contract(self) -> dict[str, Any]:
        """Return the embedding contract for diagnostics."""
        return {
            "provider": self._provider_name,
            "model": self._model_name,
            "dimensions": EMBEDDING_DIMENSION,
            "pool": self._embedding_pool.state,
        }

    def embed(self, text: str) -> list[float]:
        """Generate an embedding using the active provider.

        Args:
            text: Text to embed.

        Returns:
            list[float]: 1536-dimension embedding vector.

        Raises:
            AIGatewayValidationError: If the input text is empty.
            EmbeddingPathExhaustedError: If all embedding keys are exhausted.
            ProviderRequestError: If a non-retryable provider failure occurs.
        """
        if not text or not text.strip():
            raise AIGatewayValidationError("Embedding input text must not be empty.")

        last_error: Exception | None = None
        self._embedding_pool.reset()
        while not self._embedding_pool.is_exhausted:
            api_key = self._embedding_pool.current_key()
            
            client_kwargs = {"api_key": api_key}
            if self._base_url:
                client_kwargs["base_url"] = self._base_url
                
            client = OpenAI(**client_kwargs)
            try:
                response = client.embeddings.create(
                    input=text, 
                    model=self._model_name,
                    dimensions=EMBEDDING_DIMENSION
                )
                vector = response.data[0].embedding
                if len(vector) != EMBEDDING_DIMENSION:
                    raise ProviderRequestError(
                        f"Embedding dimension mismatch from provider response (expected {EMBEDDING_DIMENSION}, got {len(vector)})."
                    )
                logger.info(
                    "Embedding generated via AI gateway",
                    extra={
                        "provider": self._provider_name,
                        "model": self._model_name,
                        "dimensions": EMBEDDING_DIMENSION,
                    },
                )
                return vector
            except Exception as exc:  # noqa: BLE001
                last_error = exc
                _log_openai_sdk_error(
                    exc,
                    phase="embedding",
                    provider_name=self._provider_name,
                    model_name=self._model_name,
                )
                if _is_non_retryable_llm_error(exc):
                    raise ProviderRequestError(
                        f"Non-retryable embedding provider failure for '{self._provider_name}'."
                    ) from exc
                if not _is_retryable_provider_error(exc):
                    raise ProviderRequestError(
                        f"Unknown non-retryable embedding provider failure for '{self._provider_name}'."
                    ) from exc
                self._embedding_pool.rotate()

        raise EmbeddingPathExhaustedError(
            f"All keys failed for the embedding path with provider '{self._provider_name}'."
        ) from last_error


class AIGateway:
    """Single runtime control plane for gateway provider policy."""

    def __init__(self) -> None:
        """Initialize environment-driven gateway pools and provider policy.

        Raises:
            AIGatewayInitializationError: If project status is invalid or keys are missing.
        """
        settings = get_settings()
        
        project_status = settings.project_status
        if project_status not in PROFILES:
            raise AIGatewayInitializationError(
                f"PROJECT_STATUS '{project_status}' is not a valid profile. "
                f"Allowed profiles: {list(PROFILES.keys())}"
            )

        self.profile = PROFILES[project_status]
        self.routing_policy = ROUTING_POLICY[project_status]
        self.active_llm_provider = self.profile["llm_provider"]
        self.active_embedding_provider = self.profile["embedding_provider"]
        self.active_embedding_model = self.profile["embedding_model"]
        
        self.providers = PROVIDERS
        
        # Initialize all key pools explicitly
        self.openai_embedding_pool = KeyPoolManager(
            pool_name="openai_embedding_pool",
            source="OPENAI_EMBEDDING_KEYS",
            keys=settings.openai_embedding_keys,
        )
        self.openai_llm_pool = KeyPoolManager(
            pool_name="openai_llm_pool",
            source="OPENAI_LLM_KEYS",
            keys=settings.openai_llm_keys,
        )
        self.dashscope_embedding_pool = KeyPoolManager(
            pool_name="dashscope_embedding_pool",
            source="DASHSCOPE_EMBEDDING_KEYS",
            keys=settings.dashscope_embedding_keys,
        )
        self.dashscope_llm_pool = KeyPoolManager(
            pool_name="dashscope_llm_pool",
            source="DASHSCOPE_LLM_KEYS",
            keys=settings.dashscope_llm_keys,
        )

        # Validate that the active LLM provider has at least one key
        if self.active_llm_provider == "openai" and self.openai_llm_pool.is_empty:
            raise AIGatewayInitializationError(
                f"OPENAI_LLM_KEYS must contain at least one key for profile '{project_status}'."
            )
        if self.active_llm_provider == "dashscope" and self.dashscope_llm_pool.is_empty:
            raise AIGatewayInitializationError(
                f"DASHSCOPE_LLM_KEYS must contain at least one key for profile '{project_status}'."
            )
            
        # Validate that the active Embedding provider has at least one key
        if self.active_embedding_provider == "openai" and self.openai_embedding_pool.is_empty:
            raise AIGatewayInitializationError(
                f"OPENAI_EMBEDDING_KEYS must contain at least one key for profile '{project_status}'."
            )
        if self.active_embedding_provider == "dashscope" and self.dashscope_embedding_pool.is_empty:
            raise AIGatewayInitializationError(
                f"DASHSCOPE_EMBEDDING_KEYS must contain at least one key for profile '{project_status}'."
            )

    def get_langchain_model(
        self,
        task_name: str,
        model_name: str | None = None,
        stream: bool = False,
    ) -> Any:
        """Return a LangChain chat model with task-based fallbacks."""
        task_config = self._get_task_config(task_name)
        provider_name = self.active_llm_provider
        provider_config = self.providers[provider_name]
        pool = self._get_active_llm_pool()
        models = list(task_config["models"])
        params = dict(task_config.get("params", {}))

        if ChatOpenAI is None:
            raise RuntimeError("langchain_openai is not installed.")

        if model_name is not None:
            models = [model_name, *[name for name in models if name != model_name]]

        candidates: list[Any] = []
        pool.reset()
        while not pool.is_exhausted:
            api_key = pool.current_key()
            for model_name in models:
                client_kwargs: dict[str, Any] = {
                    "api_key": api_key,
                    "model": model_name,
                    "temperature": params.get("temperature", 0.0),
                    "streaming": stream,
                }
                if provider_config.get("base_url") is not None:
                    client_kwargs["base_url"] = provider_config["base_url"]
                if params.get("max_tokens") is not None:
                    client_kwargs["max_tokens"] = params["max_tokens"]

                candidates.append(ChatOpenAI(**client_kwargs))
            pool.rotate()

        if not candidates:
            raise LangChainPathExhaustedError(
                f"No LangChain candidates are available for task '{task_name}'."
            )
        if len(candidates) == 1:
            return candidates[0]
        return candidates[0].with_fallbacks(candidates[1:])

    async def aget_raw_client(
        self,
        api_key: str | None = None,
    ) -> AsyncOpenAI:
        """Return an async OpenAI-compatible client for the active provider."""
        if AsyncOpenAI is None:
            raise RuntimeError("openai.AsyncOpenAI is not installed.")

        if api_key is None:
            pool = self._get_active_llm_pool()
            if pool.is_empty:
                raise AIGatewayInitializationError(
                    f"Active provider '{self.active_llm_provider}' must contain at least one key."
                )
            pool.reset()
            api_key = pool.current_key()

        provider_config = self.providers[self.active_llm_provider]
        client_kwargs: dict[str, Any] = {"api_key": api_key}
        if provider_config.get("base_url") is not None:
            client_kwargs["base_url"] = provider_config["base_url"]
        return AsyncOpenAI(**client_kwargs)

    async def agenerate_llm_response(
        self,
        messages: list[dict[str, Any]],
        task_name: str,
        **kwargs: Any,
    ) -> Any:
        """Generate an async response using task-based routing and provider fallback."""
        if not messages:
            raise AIGatewayValidationError("LLM messages must not be empty.")

        task_config = self._get_task_config(task_name)
        provider_name = self.active_llm_provider
        pool = self._get_active_llm_pool()
        models = list(task_config["models"])
        stream = bool(kwargs.get("stream", False))
        override_model = kwargs.get("model")
        if override_model is not None:
            if override_model not in models:
                raise AIGatewayValidationError(
                    "Requested LLM model is outside the active task family."
                )
            models = [override_model]

        last_error: Exception | None = None
        pool.reset()
        while not pool.is_exhausted:
            api_key = pool.current_key()
            client = await self.aget_raw_client(api_key=api_key)
            for model_name in models:
                generation_kwargs = self._normalize_generation_kwargs(
                    {**task_config.get("params", {}), **kwargs},
                    provider_name,
                )
                generation_kwargs["model"] = model_name
                generation_kwargs["stream"] = stream
                try:
                    response = await client.chat.completions.create(
                        messages=messages,
                        **generation_kwargs,
                    )
                    logger.info(
                        "Async LLM response generated via AI gateway",
                        extra={
                            "provider": provider_name,
                            "task_name": task_name,
                            "model": model_name,
                        },
                    )
                    return response
                except Exception as exc:  # noqa: BLE001
                    last_error = exc
                    _log_openai_sdk_error(
                        exc,
                        phase="llm_chat_async",
                        provider_name=provider_name,
                        model_name=model_name,
                    )
                    if _is_non_retryable_llm_error(exc):
                        raise ProviderRequestError(
                            f"Non-retryable async LLM provider failure for '{provider_name}'."
                        ) from exc
                    if not _is_retryable_provider_error(exc):
                        raise ProviderRequestError(
                            f"Unknown non-retryable async LLM provider failure for '{provider_name}'."
                        ) from exc
                    continue
            pool.rotate()

        raise LLMPathExhaustedError(
            f"All keys and models were exhausted for task '{task_name}' on provider '{provider_name}'."
        ) from last_error

    def get_embedding_engine(self) -> EmbeddingEngine:
        """Return the embedding engine mapped to the active profile.

        Returns:
            EmbeddingEngine: Environment-driven embedding engine.
        """
        active_pool = (
            self.openai_embedding_pool
            if self.active_embedding_provider == "openai"
            else self.dashscope_embedding_pool
        )
        base_url = self.providers[self.active_embedding_provider].get("base_url")
        
        return EmbeddingEngine(
            embedding_pool=active_pool,
            provider_name=self.active_embedding_provider,
            model_name=self.active_embedding_model,
            base_url=base_url,
        )

    def get_pool_states(self) -> dict[str, KeyPoolState]:
        """Return auditable state for every explicit key pool."""
        return {
            "openai_embedding_pool": self.openai_embedding_pool.state,
            "openai_llm_pool": self.openai_llm_pool.state,
            "dashscope_embedding_pool": self.dashscope_embedding_pool.state,
            "dashscope_llm_pool": self.dashscope_llm_pool.state,
        }

    def generate_llm_response(self, messages: list[dict[str, Any]], **kwargs: Any) -> Any:
        """Generate a response using the active provider only.

        Args:
            messages: LLM messages payload.
            **kwargs: Provider-specific generation options.

        Returns:
            Any: Raw provider response object.

        Raises:
            AIGatewayValidationError: If the request payload is invalid.
            LLMPathExhaustedError: If all keys and models are exhausted.
            ProviderRequestError: If a non-retryable provider failure occurs.
        """
        if not messages:
            raise AIGatewayValidationError("LLM messages must not be empty.")

        pool = self._get_active_llm_pool()
        provider_name = self.active_llm_provider
        models = list(self.providers[provider_name]["llm_models"])
        override_model = kwargs.get("model")
        if override_model is not None:
            if override_model not in models:
                raise AIGatewayValidationError(
                    "Requested LLM model is outside the active provider family."
                )
            models = [override_model]

        last_error: Exception | None = None
        pool.reset()
        while not pool.is_exhausted:
            api_key = pool.current_key()
            for model in models:
                client = self._build_llm_client(provider_name=provider_name, api_key=api_key)
                try:
                    response = client.chat.completions.create(
                        model=model,
                        messages=messages,
                        **self._normalize_generation_kwargs(kwargs, provider_name),
                    )
                    logger.info(
                        "LLM response generated via AI gateway",
                        extra={"provider": provider_name, "model": model},
                    )
                    return response
                except Exception as exc:  # noqa: BLE001
                    last_error = exc
                    _log_openai_sdk_error(
                        exc,
                        phase="llm_chat",
                        provider_name=provider_name,
                        model_name=model,
                    )
                    logger.warning(
                        "LLM provider failed; advancing to next fallback candidate",
                        extra={
                            "provider": provider_name,
                            "model": model,
                            "error": type(exc).__name__,
                            "error_message": str(exc),
                        },
                    )
                    continue
            pool.rotate()

        raise LLMPathExhaustedError(
            f"All keys and models were exhausted for active provider '{provider_name}'."
        ) from last_error

    def _get_active_llm_pool(self) -> KeyPoolManager:
        """Return the only pool allowed for the active LLM provider."""
        if self.active_llm_provider == "openai":
            return self.openai_llm_pool
        if self.active_llm_provider == "dashscope":
            return self.dashscope_llm_pool
        raise AIGatewayInitializationError(
            "Active provider is not recognized by the gateway."
        )

    def _get_task_config(self, task_name: str) -> dict[str, Any]:
        """Return the task contract for the active profile."""
        task_config = self.routing_policy.get("tasks", {}).get(task_name)
        if task_config is None:
            raise AIGatewayValidationError(f"Task '{task_name}' is not recognized.")
        return task_config

    def _build_llm_client(self, provider_name: str, api_key: str) -> OpenAI:
        """Build a client only for the active provider path.

        Args:
            provider_name: Active provider name.
            api_key: API key selected from the active pool.

        Returns:
            OpenAI: OpenAI-compatible client for the active provider.
        """
        provider_config = self.providers[provider_name]
        base_url = provider_config.get("base_url")
        if base_url is None:
            return OpenAI(api_key=api_key)
        return OpenAI(api_key=api_key, base_url=base_url)

    @staticmethod
    def _normalize_generation_kwargs(kwargs: dict[str, Any], provider_name: str) -> dict[str, Any]:
        """Drop gateway-level keys before dispatching the provider request."""
        normalized = dict(kwargs)
        normalized.pop("model", None)
        
        if provider_name == "openai" and "max_tokens" in normalized:
            normalized["max_completion_tokens"] = normalized.pop("max_tokens")
        
        if provider_name == "dashscope" and "response_format" in normalized:
            resp_fmt = normalized["response_format"]
            if isinstance(resp_fmt, dict) and resp_fmt.get("type") == "json_schema":
                normalized["response_format"] = {"type": "json_object"}
                
        return normalized


def _is_retryable_provider_error(exc: Exception) -> bool:
    """Return whether a provider failure qualifies for fallback.

    Args:
        exc: Raised provider exception.

    Returns:
        bool: True for auth, quota, and rate-limit categories.
    """
    message = str(exc).lower()
    retryable_markers = (
        "rate limit",
        "quota",
        "insufficient_quota",
        "invalid api key",
        "unauthorized",
        "authentication",
        "401",
        "429",
    )
    return any(marker in message for marker in retryable_markers)


def _is_non_retryable_llm_error(exc: Exception) -> bool:
    """Return whether a provider failure must stop without retrying.

    Args:
        exc: Raised provider exception.

    Returns:
        bool: True for request-shape and validation categories.
    """
    message = str(exc).lower()
    non_retryable_markers = (
        "validation",
        "unsupported",
        "invalid_request_error",
        "bad request",
        "400",
        "tool schema",
        "response format",
    )
    return any(marker in message for marker in non_retryable_markers)


def _log_openai_sdk_error(
    exc: Exception,
    *,
    phase: str,
    provider_name: str,
    model_name: str,
) -> None:
    """Log raw OpenAI SDK error metadata without changing retry behavior.

    The goal is to preserve the SDK-visible failure details so request-shape
    problems can be diagnosed from logs instead of generic wrapper messages.
    """
    response = getattr(exc, "response", None)
    response_text = None
    if response is not None:
        try:
            response_text = response.text
        except Exception:  # noqa: BLE001
            response_text = "<unavailable>"

    status_code = getattr(exc, "status_code", None)
    code = getattr(exc, "code", None)
    param = getattr(exc, "param", None)
    body = getattr(exc, "body", None)
    response_status_code = getattr(response, "status_code", None)

    logger.error(
        "OpenAI SDK error details | phase=%s provider=%s model=%s exception_type=%s "
        "status_code=%s code=%s param=%s response_status_code=%s body=%r response_text=%r "
        "exception_message=%s",
        phase,
        provider_name,
        model_name,
        type(exc).__name__,
        status_code,
        code,
        param,
        response_status_code,
        body,
        response_text,
        str(exc),
    )


def create_ai_gateway() -> AIGateway:
    """Create a new AI gateway instance.

    Returns:
        AIGateway: Runtime gateway aligned to the approved Batch 1 scope.
    """
    return AIGateway()
