"""
cortex/llm/provider.py — LLMProvider ABC + Request/Response Types.

Abstraction layer over LLM API calls (OpenAI, Anthropic, Ollama, etc.).
Every Cortex module that needs LLM inference calls this interface — never
the vendor SDKs directly. This ensures:

    - Vendor independence (swap providers without changing business logic)
    - Consistent cost tracking (every call logs model, tokens, latency, cost)
    - Structured extraction via JSON schema validation
    - Embedding generation through a unified embed() method
    - Error handling + retry policy standardized across providers

Three abstract methods:
    complete() — Chat completion (text or JSON response)
    extract()  — Structured extraction (JSON conforming to a schema)
    embed()    — Text → vector embedding (for storage + retrieval)

Implementations:
    AnthropicProvider (cortex/llm/anthropic_provider.py) — Claude models
    OpenAIProvider    (cortex/llm/openai_provider.py)    — GPT-4, GPT-4.1
    OllamaProvider    (cortex/llm/ollama_provider.py)    — Local models

Convention (from ARCHITECTURE.md Module Interfaces Convention):
    - async methods for all operations
    - typed inputs/outputs via dataclasses
    - cost tracking per call (model, tokens_in, tokens_out, cost_estimate)

Dependencies: None (stdlib + abc only).
Used by: ExtractionPipeline (M2), ReconciliationEngine (M4),
         EntityResolver (M3), RetrievalPipeline (M8 agentic mode).
"""

from __future__ import annotations

import asyncio
import logging
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Retry / backoff helpers
# ---------------------------------------------------------------------------

# HTTP status codes that are safe to retry (transient server errors)
_RETRYABLE_STATUS_CODES = {429, 500, 502, 503}

# HTTP status codes that should NOT be retried (client errors, auth issues)
_NON_RETRYABLE_STATUS_CODES = {400, 401, 404}

# Default retry parameters
_MAX_RETRIES = 3
_BACKOFF_BASE_SECONDS = 1.0  # 1s, 2s, 4s exponential


def _is_retryable_exception(exc: Exception) -> bool:
    """Determine if an exception is transient and worth retrying.

    Returns True for:
        - ConnectionError, TimeoutError (network issues)
        - HTTP 429/500/502/503 errors from any SDK
    Returns False for:
        - HTTP 400/401/404 (bad request, auth, not found)
        - Any other non-transient error
    """
    # Network-level errors are always retryable
    if isinstance(exc, (ConnectionError, TimeoutError, asyncio.TimeoutError)):
        return True

    # Check for aiohttp errors (used by Ollama provider)
    exc_type_name = type(exc).__name__
    exc_module = type(exc).__module__ or ""
    if "aiohttp" in exc_module or exc_type_name in (
        "ClientError", "ClientConnectionError", "ServerTimeoutError",
        "ClientResponseError", "ServerDisconnectedError",
    ):
        status = getattr(exc, "status", None)
        if status is not None:
            if status in _NON_RETRYABLE_STATUS_CODES:
                return False
            if status in _RETRYABLE_STATUS_CODES:
                return True
        # Generic connection errors are retryable
        return True

    # Check for HTTP status code on SDK exceptions (openai, anthropic, httpx)
    status_code = (
        getattr(exc, "status_code", None)
        or getattr(exc, "status", None)
        or getattr(exc, "code", None)
    )

    if status_code is not None:
        try:
            code = int(status_code)
        except (ValueError, TypeError):
            return False
        if code in _NON_RETRYABLE_STATUS_CODES:
            return False
        if code in _RETRYABLE_STATUS_CODES:
            return True

    # Check for "rate limit" in the error message as a fallback
    msg = str(exc).lower()
    if "rate limit" in msg or "429" in msg:
        return True
    if "timeout" in msg or "connection" in msg:
        return True

    return False


async def retry_with_backoff(
    coro_factory,
    *,
    max_retries: int = _MAX_RETRIES,
    backoff_base: float = _BACKOFF_BASE_SECONDS,
    operation_name: str = "llm_call",
) -> Any:
    """Execute an async operation with exponential backoff retry.

    Args:
        coro_factory: A zero-arg callable that returns a new coroutine each time.
        max_retries: Maximum number of retry attempts (default 3).
        backoff_base: Base delay in seconds (doubles each retry: 1s, 2s, 4s).
        operation_name: Name for logging purposes.

    Returns:
        The result of the successful coroutine call.

    Raises:
        The last exception if all retries are exhausted.
    """
    last_exc: Exception = RuntimeError(f"No attempts made for {operation_name}")
    for attempt in range(max_retries + 1):  # attempt 0 = first try, 1..N = retries
        try:
            return await coro_factory()
        except Exception as exc:
            last_exc = exc
            if attempt >= max_retries:
                logger.warning(
                    "%s: all %d retries exhausted. Last error: %s",
                    operation_name, max_retries, exc,
                )
                raise
            if not _is_retryable_exception(exc):
                logger.debug(
                    "%s: non-retryable error (attempt %d): %s",
                    operation_name, attempt + 1, exc,
                )
                raise
            delay = backoff_base * (2 ** attempt)
            logger.info(
                "%s: retryable error (attempt %d/%d): %s — backing off %.1fs",
                operation_name, attempt + 1, max_retries + 1, exc, delay,
            )
            await asyncio.sleep(delay)
    raise last_exc  # unreachable but satisfies type checker


# ---------------------------------------------------------------------------
# Request / Response dataclasses
# ---------------------------------------------------------------------------

@dataclass
class CompletionRequest:
    """Input to LLMProvider.complete().

    Attributes:
        system: System prompt text (agent instructions).
        user: User message text (the query/conversation).
        model: Model identifier (e.g., 'gpt-4.1-nano-2025-04-14').
        response_format: 'text' for free-form, 'json' for JSON-only output.
        max_tokens: Optional cap on response length.
        temperature: Sampling temperature (0.0 = deterministic).
    """
    system: str
    user: str
    model: str
    response_format: str = "text"   # "text" | "json"
    max_tokens: Optional[int] = None
    temperature: float = 0.0


@dataclass
class CompletionResponse:
    """Output of LLMProvider.complete().

    Attributes:
        text: The generated text response.
        model: Which model actually produced the response (may differ from request
               if the provider substituted a fallback model).
        tokens_in: Input tokens consumed (for cost tracking).
        tokens_out: Output tokens generated.
        cost_estimate: Estimated cost in USD (best-effort, based on token pricing).
        latency_ms: Wall-clock time for the API call in milliseconds.
    """
    text: str
    model: str
    tokens_in: int
    tokens_out: int
    cost_estimate: float            # USD, best-effort
    latency_ms: float


@dataclass
class ExtractionRequest:
    """Input to LLMProvider.extract() — structured JSON extraction.

    Like CompletionRequest but with a JSON schema that the output must conform to.
    Used by ReconciliationEngine for structured classification responses.

    Attributes:
        system: System prompt text.
        user: User message text.
        schema: JSON Schema dict describing the expected output structure.
        model: Model identifier.
        max_tokens: Optional cap on response length.
        temperature: Sampling temperature (0.0 = deterministic).
    """
    system: str
    user: str
    schema: dict[str, Any]          # JSON schema describing expected output
    model: str
    max_tokens: Optional[int] = None
    temperature: float = 0.0


@dataclass
class ExtractionResponse:
    """Output of LLMProvider.extract()."""
    data: dict[str, Any]            # parsed JSON matching the schema
    model: str
    tokens_in: int
    tokens_out: int
    cost_estimate: float
    latency_ms: float


@dataclass
class EmbedRequest:
    """Input to LLMProvider.embed()."""
    texts: list[str]
    model: str
    input_type: str = "query"       # "query" | "document" (Voyage convention)


@dataclass
class EmbedResponse:
    """Output of LLMProvider.embed()."""
    embeddings: list[list[float]]
    model: str
    tokens_in: int                  # total tokens across all texts
    tokens_out: int = 0             # always 0 for embeddings
    cost_estimate: float = 0.0
    latency_ms: float = 0.0


@dataclass
class UsageRecord:
    """Single cost-tracking record written after every LLM call."""
    provider: str
    model: str
    call_type: str                  # "complete" | "extract" | "embed"
    tokens_in: int
    tokens_out: int
    cost_estimate: float
    latency_ms: float
    timestamp: float = field(default_factory=time.time)


# ---------------------------------------------------------------------------
# Abstract base class
# ---------------------------------------------------------------------------

class LLMProvider(ABC):
    """
    Abstraction over LLM API calls.

    Implementations: AnthropicProvider, OpenAIProvider, OllamaProvider.

    All methods are async. Every call logs a UsageRecord for cost tracking.
    API keys are read from environment variables — never hardcode them here.
    """

    # Each subclass declares its own default models:
    DEFAULT_COMPLETION_MODEL: str = ""
    DEFAULT_EMBED_MODEL: str = ""

    def __init__(self) -> None:
        self._usage_log: list[UsageRecord] = []

    # --- core methods -------------------------------------------------------

    @abstractmethod
    async def complete(
        self,
        system: str,
        user: str,
        model: str | None = None,
        response_format: str = "text",
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> CompletionResponse:
        """
        Send a system+user prompt and return the completion.

        Args:
            system: System prompt.
            user: User message.
            model: Override default model. If None, uses DEFAULT_COMPLETION_MODEL.
            response_format: "text" (default) or "json".
            max_tokens: Hard cap on output tokens.
            temperature: Sampling temperature (0 = deterministic).

        Returns:
            CompletionResponse with text + cost metrics.
        """

    @abstractmethod
    async def extract(
        self,
        system: str,
        user: str,
        schema: dict[str, Any],
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """
        Structured JSON extraction — forces output to match *schema*.

        Args:
            system: System prompt (should describe the extraction task).
            user: User message or document to extract from.
            schema: JSON schema dict describing expected output structure.
            model: Override default model.
            max_tokens: Hard cap on output tokens.
            temperature: Sampling temperature.

        Returns:
            ExtractionResponse with parsed dict + cost metrics.
        """

    @abstractmethod
    async def embed(
        self,
        texts: list[str],
        model: str | None = None,
        input_type: str = "query",
    ) -> EmbedResponse:
        """
        Generate embeddings for a list of texts.

        Args:
            texts: List of strings to embed.
            model: Override default embed model.
            input_type: "query" or "document" (used by Voyage).

        Returns:
            EmbedResponse with list-of-list-of-floats + cost metrics.
        """

    # --- helpers ------------------------------------------------------------

    def _log_usage(self, record: UsageRecord) -> None:
        """Append to in-memory usage log and emit a structured log line."""
        self._usage_log.append(record)
        logger.info(
            "llm_cost provider=%s model=%s call=%s "
            "tokens_in=%d tokens_out=%d cost_usd=%.6f latency_ms=%.1f",
            record.provider,
            record.model,
            record.call_type,
            record.tokens_in,
            record.tokens_out,
            record.cost_estimate,
            record.latency_ms,
        )

    def usage_summary(self) -> dict[str, Any]:
        """Return aggregated usage stats across all calls this session."""
        total_cost = sum(r.cost_estimate for r in self._usage_log)
        total_in = sum(r.tokens_in for r in self._usage_log)
        total_out = sum(r.tokens_out for r in self._usage_log)
        return {
            "calls": len(self._usage_log),
            "tokens_in": total_in,
            "tokens_out": total_out,
            "cost_usd": round(total_cost, 6),
        }

    def _resolve_model(self, model: str | None, default: str) -> str:
        return model or default or self.DEFAULT_COMPLETION_MODEL
