"""Configurable LLM client — supports LiteLLM (100+ providers) or mock/stub mode.

Reads AI_MODEL, AI_API_KEY, AI_PROVIDER from environment. If not set, uses mock mode
which returns predefined actions based on keyword matching. This allows
tests to run without external API dependencies.

LiteLLM provides a unified interface to OpenAI, Anthropic, Google, Azure,
AWS Bedrock, Ollama, and many more providers.

Generic functions:
- ``llm_complete()`` — generic chat completion with retry, cost tracking
- ``llm_embed()`` — generic text embedding
- ``get_api_credentials()`` — centralised credential/provider lookup
- ``build_model()`` — centralised LiteLLM model string builder
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import uuid
from typing import Any, TYPE_CHECKING

import litellm

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

# ──────────────────────────────────────────────────────────────────────────
# Constants
# ──────────────────────────────────────────────────────────────────────────

MAX_INPUT_CHARS = 8000

# OpenRouter for embeddings (Ollama Cloud has no embedding endpoint)
OPENROUTER_API_KEY = os.environ.get("API_KEY_OPENROUTER", "")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
OPENROUTER_EMBEDDING_MODEL = os.environ.get(
    "SEARCH_EMBEDDING_MODEL", "openai/text-embedding-3-small"
)
EMBEDDING_DIMENSIONS = 768  # Must match DB column vector(768)

# Default retry settings
DEFAULT_TIMEOUT = 30
DEFAULT_MAX_RETRIES = 2
BASE_BACKOFF_SECONDS = 1.0

# Transient error keywords for retry classification
_TRANSIENT_KEYWORDS = frozenset(
    {
        "timeout",
        "timed out",
        "rate limit",
        "rate_limit",
        "429",
        "503",
        "502",
        "504",
        "service unavailable",
        "overloaded",
        "connection reset",
        "connection aborted",
        "temporary",
    }
)

# Permanent error keywords — fail immediately, no retry
_PERMANENT_KEYWORDS = frozenset(
    {
        "authentication",
        "auth",
        "401",
        "403",
        "unauthorized",
        "forbidden",
        "invalid api key",
        "invalid_api_key",
        "validation",
        "invalid_request",
        "400",
        "bad request",
        "model_not_found",
        "not found",
    }
)


# ──────────────────────────────────────────────────────────────────────────
# Centralised helper functions
# ──────────────────────────────────────────────────────────────────────────


async def get_api_credentials(
    db: AsyncSession | None,
    tenant_id: uuid.UUID | None,
) -> tuple[str | None, str | None, str | None]:
    """Get API key, base_url and provider_type for LLM/embedding calls.

    Priority:
    1. OpenRouter env var (API_KEY_OPENROUTER) – dedicated embedding provider
    2. Default AI provider from DB (fallback)
    3. API_KEY_OLLAMA_CLOUD env var (last resort)

    Args:
        db: Optional async DB session for provider lookup.
        tenant_id: Optional tenant ID for provider lookup.

    Returns:
        Tuple of (api_key, api_base, provider_type) — any may be ``None``.
    """
    # OpenRouter is the primary embedding provider
    if OPENROUTER_API_KEY:
        return OPENROUTER_API_KEY, OPENROUTER_BASE_URL, "openai"

    # Fallback to DB provider
    if db and tenant_id:
        try:
            from app.plugins.builtins.ai_assistant.contracts import get_default_provider

            provider = await get_default_provider(db, tenant_id)
            if provider and provider.api_key:
                return provider.api_key, provider.base_url, provider.provider_type
        except Exception:
            logger.debug("Failed to get provider from DB, falling back to env")

    env_key = os.environ.get("API_KEY_OLLAMA_CLOUD", "")
    return (env_key if env_key else None), None, None


def build_model(model: str, provider_type: str | None) -> str:
    """Build LiteLLM model string with provider prefix.

    If ``provider_type`` is given, strips any existing prefix from ``model``
    and prepends ``provider_type``.

    Args:
        model: Model name, optionally already prefixed (e.g. ``openai/gpt-4o``).
        provider_type: Provider prefix to apply (e.g. ``openai``, ``anthropic``).

    Returns:
        LiteLLM-compatible model string (e.g. ``openai/gpt-4o``).
    """
    if provider_type:
        model_parts = model.split("/", 1)
        return f"{provider_type}/{model_parts[-1]}"
    return model


def _classify_error(exc: Exception) -> str:
    """Classify an exception as ``transient`` or ``permanent``.

    Uses string matching on the exception message/type name against known
    patterns. Falls back to ``transient`` for unknown errors (safer to retry).

    Args:
        exc: The exception to classify.

    Returns:
        ``"transient"`` or ``"permanent"``.
    """
    msg = str(exc).lower()
    exc_type_name = type(exc).__name__.lower()

    # Check permanent first — auth errors should never be retried
    if any(kw in msg or kw in exc_type_name for kw in _PERMANENT_KEYWORDS):
        return "permanent"
    if any(kw in msg or kw in exc_type_name for kw in _TRANSIENT_KEYWORDS):
        return "transient"
    # asyncio.TimeoutError is always transient
    if isinstance(exc, (asyncio.TimeoutError, TimeoutError)):
        return "transient"
    # Default: treat as transient (safe to retry)
    return "transient"


def _extract_cost_usd(response: Any, model: str) -> float:
    """Extract cost in USD from a LiteLLM response.

    Uses ``litellm.completion_cost`` when available, otherwise returns 0.0.

    Args:
        response: LiteLLM response object.
        model: Model string used for the call.

    Returns:
        Estimated cost in USD, or 0.0 if unavailable.
    """
    try:
        cost = litellm.completion_cost(response)
        if cost is not None:
            return float(cost)
    except Exception:
        logger.debug("litellm.completion_cost failed, using fallback")
    return 0.0


def _extract_usage(response: Any) -> dict[str, int]:
    """Extract token usage from a LiteLLM response.

    Args:
        response: LiteLLM response object.

    Returns:
        Dict with ``prompt_tokens``, ``completion_tokens``, ``total_tokens``.
    """
    usage = getattr(response, "usage", None)
    if usage is None:
        return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
    completion_tokens = getattr(usage, "completion_tokens", 0) or 0
    total_tokens = getattr(usage, "total_tokens", 0) or (prompt_tokens + completion_tokens)
    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": total_tokens,
    }


# ──────────────────────────────────────────────────────────────────────────
# Generic LLM functions
# ──────────────────────────────────────────────────────────────────────────


async def llm_complete(
    model: str,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]] | None = None,
    temperature: float = 0.3,
    max_tokens: int = 1000,
    api_key: str | None = None,
    api_base: str | None = None,
    provider: str | None = None,
    response_format: dict[str, Any] | None = None,
    timeout: int = DEFAULT_TIMEOUT,
    max_retries: int = DEFAULT_MAX_RETRIES,
) -> dict[str, Any]:
    """Generic LLM chat completion via LiteLLM with retry and cost tracking.

    Supports 100+ providers through LiteLLM's unified interface.
    Transient errors (timeout, rate-limit) are retried with exponential
    backoff. Permanent errors (auth, validation) fail immediately.

    Args:
        model: Model name (e.g. ``gpt-4o``, ``openai/gpt-4o``).
        messages: Chat messages list (``[{"role": ..., "content": ...}]``).
        tools: Optional list of tool/function definitions.
        temperature: Sampling temperature (default 0.3).
        max_tokens: Maximum tokens to generate (default 1000).
        api_key: Override API key. If ``None``, uses env/DB lookup.
        api_base: Override API base URL.
        provider: Provider prefix (e.g. ``openai``, ``anthropic``).
        response_format: Optional response format spec (e.g. JSON mode).
        timeout: Request timeout in seconds (default 30).
        max_retries: Max retry attempts for transient errors (default 2).

    Returns:
        Dict with keys: ``content``, ``usage``, ``cost_usd``, ``model``,
        ``raw_response`` (the LiteLLM response object for advanced use).

    Raises:
        Exception: Permanent errors or after exhausting retries.
    """
    # Build LiteLLM model string
    litellm_model = build_model(model, provider)

    # Build kwargs
    kwargs: dict[str, Any] = {
        "model": litellm_model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "timeout": timeout,
    }
    if api_key:
        kwargs["api_key"] = api_key
    if api_base:
        kwargs["api_base"] = api_base
    if tools:
        kwargs["tools"] = tools
    if response_format:
        kwargs["response_format"] = response_format

    last_exc: Exception | None = None

    for attempt in range(max_retries + 1):
        try:
            response = await litellm.acompletion(**kwargs)
            content = response.choices[0].message.content or ""
            usage = _extract_usage(response)
            cost_usd = _extract_cost_usd(response, litellm_model)

            logger.debug(
                "llm_complete success: model=%s tokens=%d cost=$%.6f attempt=%d",
                litellm_model,
                usage["total_tokens"],
                cost_usd,
                attempt + 1,
            )
            return {
                "content": content,
                "usage": usage,
                "cost_usd": cost_usd,
                "model": litellm_model,
                "raw_response": response,
            }
        except Exception as exc:
            last_exc = exc
            error_class = _classify_error(exc)

            if error_class == "permanent" or attempt >= max_retries:
                logger.error(
                    "llm_complete failed (permanent/exhausted): model=%s attempt=%d error=%s",
                    litellm_model,
                    attempt + 1,
                    exc,
                )
                raise

            # Transient error — retry with exponential backoff
            backoff = BASE_BACKOFF_SECONDS * (2**attempt)
            logger.warning(
                "llm_complete transient error (attempt %d/%d), retrying in %.1fs: %s",
                attempt + 1,
                max_retries + 1,
                backoff,
                exc,
            )
            await asyncio.sleep(backoff)

    # Should not reach here, but satisfy type checker
    assert last_exc is not None
    raise last_exc


async def llm_embed(
    texts: str | list[str],
    model: str | None = None,
    db: AsyncSession | None = None,
    tenant_id: uuid.UUID | None = None,
    api_key: str | None = None,
    api_base: str | None = None,
    provider: str | None = None,
    dimensions: int | None = None,
    timeout: int = DEFAULT_TIMEOUT,
) -> list[list[float]]:
    """Generic text embedding via LiteLLM.

    Handles both single-text and batch embedding. Uses centralised
    credential lookup when ``api_key`` is not provided.

    Args:
        texts: Single text string or list of texts to embed.
        model: Embedding model name (default: ``openai/text-embedding-3-small``).
        db: Optional DB session for API key lookup.
        tenant_id: Optional tenant ID for API key lookup.
        api_key: Override API key. If ``None``, uses env/DB lookup.
        api_base: Override API base URL.
        provider: Provider prefix override.
        dimensions: Override embedding dimensions.
        timeout: Request timeout in seconds (default 30).

    Returns:
        List of embedding vectors (each a list of floats). For a single
        text input, returns a one-element list.
    """
    # Normalise to list input
    single_input = isinstance(texts, str)
    text_list = [texts] if single_input else texts

    if not text_list:
        return []

    # Truncate inputs
    truncated = [t[:MAX_INPUT_CHARS] for t in text_list]

    # Resolve model
    embedding_model = model or OPENROUTER_EMBEDDING_MODEL

    # Resolve credentials
    if not api_key:
        resolved_key, resolved_base, resolved_provider = await get_api_credentials(
            db, tenant_id
        )
        api_key = resolved_key
        if not api_base:
            api_base = resolved_base
        if not provider:
            provider = resolved_provider

    # Build LiteLLM model string
    litellm_model = build_model(embedding_model, provider)

    # Build kwargs
    litellm_kwargs: dict[str, Any] = {
        "model": litellm_model,
        "input": truncated[0] if single_input else truncated,
        "timeout": timeout,
    }
    if api_key:
        litellm_kwargs["api_key"] = api_key
    if api_base:
        litellm_kwargs["api_base"] = api_base

    # Request specific dimensions for text-embedding-3 models
    effective_dims = dimensions or EMBEDDING_DIMENSIONS
    if "text-embedding-3" in litellm_model:
        litellm_kwargs["dimensions"] = effective_dims

    try:
        response = await litellm.aembedding(**litellm_kwargs)
        embeddings = [d["embedding"] for d in response.data]
        logger.debug(
            "llm_embed success: model=%s count=%d dims=%d",
            litellm_model,
            len(embeddings),
            len(embeddings[0]) if embeddings else 0,
        )
        return embeddings
    except Exception:
        logger.warning("llm_embed failed: model=%s", litellm_model, exc_info=True)
        return [[] for _ in text_list]


# ──────────────────────────────────────────────────────────────────────────
# LLMClient class (AI Copilot — backward compatible)
# ──────────────────────────────────────────────────────────────────────────


class LLMResponse:
    """Structured LLM response containing proposed actions."""

    def __init__(
        self, message: str, proposed_actions: list[dict[str, Any]], confidence: float = 0.8
    ):
        self.message = message
        self.proposed_actions = proposed_actions
        self.confidence = confidence

    def to_dict(self) -> dict[str, Any]:
        return {
            "message": self.message,
            "proposed_actions": self.proposed_actions,
            "confidence": self.confidence,
        }


class LLMClient:
    """LLM client that translates natural language to proposed API actions.

    Modes:
    - If AI_MODEL and AI_API_KEY are set: calls LiteLLM chat completions API
    - Otherwise: mock/stub mode with keyword-based action mapping

    LiteLLM model format: "provider/model_name" (e.g. "openai/gpt-4o", "anthropic/claude-3-sonnet", "ollama/llama3")
    """

    def __init__(
        self, model: str | None = None, api_key: str | None = None, api_base: str | None = None,
        provider: str | None = None,
    ) -> None:
        self.model = model or os.environ.get("AI_MODEL", "")
        self.api_key = api_key or os.environ.get("AI_API_KEY", "")
        self.api_base = api_base or os.environ.get("AI_API_BASE", "")
        self.provider = provider or os.environ.get("AI_PROVIDER", "openai")
        self.is_mock = not bool(self.model and self.api_key)

    async def generate(self, user_query: str, context: dict[str, Any] | None = None) -> LLMResponse:
        """Generate proposed actions from natural language query.

        Args:
            user_query: Natural language input from user
            context: Optional context (e.g. current page, selected entity)

        Returns:
            LLMResponse with message and proposed_actions list
        """
        if self.is_mock:
            return await self._mock_generate(user_query, context or {})
        return await self._api_generate(user_query, context or {})

    async def _mock_generate(self, query: str, context: dict[str, Any]) -> LLMResponse:
        """Mock/stub mode — keyword-based action mapping for tests."""
        from app.ai.action_mapper import map_query_to_actions

        actions = map_query_to_actions(query, context)
        if actions:
            return LLMResponse(
                message=f"I found {len(actions)} possible action(s) based on your request.",
                proposed_actions=actions,
                confidence=0.85,
            )
        return LLMResponse(
            message="I couldn't determine a specific action from your request. Could you be more specific?",
            proposed_actions=[],
            confidence=0.3,
        )

    async def _api_generate(self, query: str, context: dict[str, Any]) -> LLMResponse:
        """Call LLM via ``llm_complete()`` (delegates to LiteLLM).

        Supports 100+ providers through a single API:
        - OpenAI: "openai/gpt-4o"
        - Anthropic: "anthropic/claude-3-sonnet"
        - Google: "gemini/gemini-pro"
        - Azure: "azure/<deployment-name>"
        - Ollama: "ollama/llama3"
        - And many more.
        """
        system_prompt = self._build_system_prompt(context)
        user_prompt = f"User request: {query}\n\nRespond with proposed actions as JSON."

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

        try:
            result = await llm_complete(
                model=self.model,
                messages=messages,
                temperature=0.3,
                max_tokens=1000,
                api_key=self.api_key or None,
                api_base=self.api_base or None,
                provider=self.provider,
            )
            return self._parse_llm_response(result["content"])
        except Exception as e:
            logger.error("LLM API call failed: %s", e)
            # Fall back to mock mode on API error
            return LLMResponse(
                message=f"LLM API call failed: {e}. Falling back to keyword matching.",
                proposed_actions=[],
                confidence=0.1,
            )

    def _build_system_prompt(self, context: dict[str, Any]) -> str:
        """Build system prompt describing available API actions."""
        available_apis = [
            {"method": "GET", "path": "/api/v1/contacts", "description": "List contacts (persons and companies)"},
            {"method": "POST", "path": "/api/v1/contacts", "description": "Create a contact (person or company)"},
            {
                "method": "GET",
                "path": "/api/v1/contacts/{id}",
                "description": "Get contact details",
            },
            {
                "method": "PATCH",
                "path": "/api/v1/contacts/{id}",
                "description": "Update a contact",
            },
            {
                "method": "DELETE",
                "path": "/api/v1/contacts/{id}",
                "description": "Delete a contact",
            },
            {"method": "GET", "path": "/api/v1/workflows", "description": "List workflows"},
            {"method": "POST", "path": "/api/v1/workflows", "description": "Create a workflow"},
            {"method": "GET", "path": "/api/v1/calendar/entries", "description": "List calendar entries"},
            {"method": "POST", "path": "/api/v1/calendar/entries", "description": "Create a calendar entry"},
            {"method": "GET", "path": "/api/v1/dms/files", "description": "List DMS files"},
        ]
        context_str = json.dumps(context) if context else "{}"
        return (
            "You are an AI copilot for LeoCRM. Based on the user's natural language request, "
            "propose one or more API actions. Always respond with a JSON object containing: "
            '"message": a human-readable summary, '
            '"proposed_actions": an array of {method, path, body, description, confidence}. '
            f"Available API endpoints: {json.dumps(available_apis)}. "
            f"Current context: {context_str}. "
            "Never execute actions directly — only propose them for user confirmation."
        )

    def _parse_llm_response(self, content: str) -> LLMResponse:
        """Parse LLM JSON response into LLMResponse."""
        try:
            parsed = json.loads(content)
            return LLMResponse(
                message=parsed.get("message", "Here are the proposed actions."),
                proposed_actions=parsed.get("proposed_actions", []),
                confidence=parsed.get("confidence", 0.8),
            )
        except (json.JSONDecodeError, KeyError):
            logger.warning("Failed to parse LLM response as JSON: %s", content[:200])
            return LLMResponse(
                message=content,
                proposed_actions=[],
                confidence=0.3,
            )


# Global client instance
_client: LLMClient | None = None


def get_llm_client() -> LLMClient:
    """Get or create the global LLM client instance."""
    global _client
    if _client is None:
        _client = LLMClient()
    return _client


def reset_llm_client() -> None:
    """Reset the global client (for testing)."""
    global _client
    _client = None