"""Base LLM layer for Fluxibly.

Defines BaseLLM (abstract), LLMConfig, ReasoningConfig, and PricingConfig.
All LLM implementations extend BaseLLM and implement prepare() + forward().
"""

from __future__ import annotations

from collections.abc import Generator
from typing import Any, Literal

from pydantic import BaseModel, Field

from fluxibly.logging import Logger
from fluxibly.schema.response import LLMResponse


class ReasoningConfig(BaseModel):
    """Configuration for reasoning/thinking models."""

    enabled: bool = False
    effort: Literal["low", "medium", "high"] = Field(
        default="medium", description="Reasoning effort level"
    )
    budget_tokens: int | None = Field(
        default=None, description="Max reasoning tokens (Anthropic)"
    )


class PricingConfig(BaseModel):
    """Optional per-model pricing for cost tracking.

    All prices are in USD per 1M tokens. When set, TokenUsage will
    automatically calculate cost after each forward() call.
    """

    input_price: float = Field(..., description="USD per 1M input tokens")
    output_price: float = Field(..., description="USD per 1M output tokens")
    cached_input_price: float | None = Field(
        default=None, description="USD per 1M cached input tokens"
    )
    reasoning_price: float | None = Field(
        default=None, description="USD per 1M reasoning tokens"
    )


class LLMConfig(BaseModel):
    """Configuration for LLM initialization."""

    # --- Model ---
    model: str = Field(
        ...,
        description="Model identifier, e.g. 'gpt-4o', 'claude-sonnet-4-20250514'",
    )
    api_key: str | None = Field(
        default=None, description="API key (or use env var)"
    )
    api_base: str | None = Field(default=None, description="Custom endpoint")

    # --- Generation Parameters ---
    temperature: float | None = None
    max_output_tokens: int | None = Field(
        default=None, description="Max tokens to generate"
    )
    max_input_tokens: int | None = Field(
        default=None, description="Max input context tokens"
    )
    top_p: float | None = None
    frequency_penalty: float | None = None
    presence_penalty: float | None = None
    stop: list[str] | None = Field(default=None, description="Stop sequences")

    # --- Reasoning (for models that support it) ---
    reasoning: ReasoningConfig | None = None

    # --- Request Settings ---
    timeout: int = Field(default=60, description="Request timeout in seconds")
    max_retries: int = Field(
        default=3, description="Max retries on transient errors"
    )
    retry_backoff: float = Field(
        default=2.0, description="Exponential backoff multiplier"
    )
    streaming: bool = False

    # --- Response API vs Chat API (OpenAI-specific) ---
    api_type: Literal["responses", "chat"] = Field(
        default="responses", description="OpenAI API type"
    )

    # --- Pricing (optional, for cost tracking) ---
    pricing: PricingConfig | None = None

    # --- Logging ---
    log_level: str = Field(
        default="INFO", description="'DEBUG' | 'INFO' | 'WARNING' | 'ERROR'"
    )

    # --- Additional provider-specific params ---
    additional_params: dict[str, Any] = Field(default_factory=dict)


class BaseLLM:
    """Abstract base class for all LLM implementations.

    Does NOT implement prepare() or forward() — extended classes do.
    Provides shared infrastructure: config, logging, message/tool management.

    Each forward() call = exactly one LLM inference.
    """

    def __init__(self, config: LLMConfig) -> None:
        self.config = config
        self.messages: list[dict[str, Any]] = []
        self.tools: list[dict[str, Any]] = []
        self.skills: list[dict[str, Any]] = []
        self.logger = Logger(
            component="llm", level=config.log_level, model=config.model
        )

    # --- Message Management ---

    def set_messages(self, messages: list[dict[str, Any]]) -> None:
        """Set conversation history. Each message: {role, content, ...}"""
        self.messages = messages

    def add_message(self, message: dict[str, Any]) -> None:
        """Append a single message to conversation history."""
        self.messages.append(message)

    def get_messages(self) -> list[dict[str, Any]]:
        """Return current conversation history."""
        return self.messages

    def clear_messages(self) -> None:
        """Clear conversation history."""
        self.messages = []

    # --- Tool Management ---

    def set_tools(self, tools: list[dict[str, Any]]) -> None:
        """Set available tools. Format depends on provider."""
        self.tools = tools

    # --- Skill Management ---

    def set_skills(self, skills: list[dict[str, Any]]) -> None:
        """Set Agent Skills. Skills are instruction/resource bundles (Agent Skills standard).

        Each skill: {skill_id, source, version, ...}
        Supported by OpenAI (shell environment) and Anthropic (container).
        """
        self.skills = skills

    # --- Core Methods (abstract) ---

    def prepare(self) -> dict[str, Any]:
        """Prepare all assets into provider-specific format for inference.

        Returns:
            dict with provider-specific formatted inputs.

        Raises:
            NotImplementedError: Must be implemented by subclass.
        """
        raise NotImplementedError

    def forward(self, **kwargs: Any) -> LLMResponse:
        """Execute a single LLM inference with prepared inputs.

        Returns:
            LLMResponse with standardized output format.

        Raises:
            NotImplementedError: Must be implemented by subclass.
        """
        raise NotImplementedError

    def forward_stream(self, **kwargs: Any) -> Generator[Any, None, None]:
        """Streaming variant of forward().

        Yields:
            LLMResponseChunk objects.

        Raises:
            NotImplementedError: Streaming is not yet supported for this provider.
        """
        raise NotImplementedError(
            f"Streaming is not yet supported for {self.__class__.__name__}. "
            f"Use forward() for non-streaming inference."
        )

    # --- Shared Helpers ---

    def _apply_pricing(self, response: LLMResponse) -> None:
        """Apply pricing calculations to the response if configured."""
        if self.config.pricing and response.metadata.usage:
            response.metadata.usage.calculate_cost(self.config.pricing)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(model={self.config.model})"
