"""Unified LLM module for all model definitions, clients, and API calls."""

from __future__ import annotations

import logging
import os
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal

from anthropic import AsyncAnthropicFoundry
from openai import AsyncAzureOpenAI
from pydantic.dataclasses import dataclass as pydantic_dataclass

if TYPE_CHECKING:
    from anthropic.types import Message as AnthropicMessage
    from openai.types.chat import ChatCompletion, ChatCompletionMessageParam

logger = logging.getLogger(__name__)


# =============================================================================
# Model Definitions
# =============================================================================

# Pricing is in USD per 1M tokens. See:
# https://docs.anthropic.com/en/docs/about-claude/pricing
# https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/


@pydantic_dataclass
class LLM:
    """Base LLM configuration with pricing info."""

    name: str  # On Azure OpenAI Service, this is the deployment name
    max_tokens: int
    model_type: Literal["openai", "anthropic", "google"]
    input_cost: float | None = None  # USD per 1M tokens
    cached_input_cost: float | None = None  # USD per 1M tokens (cached input)
    output_cost: float | None = None  # USD per 1M tokens


@pydantic_dataclass
class OpenAI_GPT_4_1(LLM):
    """OpenAI GPT-4.1 model."""

    name: str = "gpt-4.1"
    model_type: Literal["openai", "anthropic", "google"] = "openai"
    max_tokens: int = 100000
    input_cost: float = 2.00
    cached_input_cost: float = 0.50
    output_cost: float = 8.00


@pydantic_dataclass
class OpenAI_GPT_5_Mini(LLM):
    """OpenAI GPT-5-mini model."""

    name: str = "gpt-5-mini"
    model_type: Literal["openai", "anthropic", "google"] = "openai"
    max_tokens: int = 200000
    input_cost: float = 0.25
    cached_input_cost: float = 0.03
    output_cost: float = 2.00


@pydantic_dataclass
class Anthropic_Claude_Sonnet_4_5(LLM):
    """Anthropic Claude 4.5 Sonnet via Azure Foundry."""

    name: str = "claude-sonnet-4-5"
    model_type: Literal["openai", "anthropic", "google"] = "anthropic"
    max_tokens: int = 200000
    input_cost: float = 3.00
    output_cost: float = 15.00


@pydantic_dataclass
class Anthropic_Claude_Haiku_4_5(LLM):
    """Anthropic Claude 4.5 Haiku via Azure Foundry."""

    name: str = "claude-haiku-4-5"
    model_type: Literal["openai", "anthropic", "google"] = "anthropic"
    max_tokens: int = 200000
    input_cost: float = 1.00
    output_cost: float = 5.00


# =============================================================================
# LLM Client Setup
# =============================================================================

# Read environment variables once at module load
_AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY")
_ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_FOUNDRY_API_KEY")
_ANTHROPIC_BASE_URL = os.environ.get("ANTHROPIC_FOUNDRY_BASE_URL")

# Clear ANTHROPIC_FOUNDRY_RESOURCE to prevent SDK conflict (mutually exclusive with base_url)
os.environ.pop("ANTHROPIC_FOUNDRY_RESOURCE", None)


def _create_openai_client() -> AsyncAzureOpenAI | None:
    if _AZURE_OPENAI_API_KEY:
        return AsyncAzureOpenAI()  # SDK auto-reads AZURE_OPENAI_* and OPENAI_API_VERSION
    return None


def _create_anthropic_client() -> AsyncAnthropicFoundry | None:
    if _ANTHROPIC_API_KEY and _ANTHROPIC_BASE_URL:
        return AsyncAnthropicFoundry(api_key=_ANTHROPIC_API_KEY, base_url=_ANTHROPIC_BASE_URL)
    return None


def get_llm_client(model_type: str) -> AsyncAzureOpenAI | AsyncAnthropicFoundry | None:
    """Get a fresh LLM client for the request.

    Creates a new client for each request to avoid event loop issues
    with Django dev server.
    """
    if model_type == "openai":
        return _create_openai_client()
    if model_type == "anthropic":
        return _create_anthropic_client()
    return None


# Keep module-level clients for backwards compatibility
_openai_client = _create_openai_client()
_anthropic_client = _create_anthropic_client()

llm_clients: dict[str, AsyncAzureOpenAI | AsyncAnthropicFoundry | None] = {
    "openai": _openai_client,
    "anthropic": _anthropic_client,
}


# =============================================================================
# Response Types
# =============================================================================


@dataclass
class LLMUsage:
    """Unified usage stats for both OpenAI and Anthropic responses."""

    input_tokens: int
    output_tokens: int


@dataclass
class LLMResponse:
    """Unified response wrapper for both OpenAI and Anthropic API responses."""

    content: str
    usage: LLMUsage
    raw_response: ChatCompletion | AnthropicMessage


# =============================================================================
# LLM API Call
# =============================================================================


async def call_llm(
    llm: LLM,
    messages: list[ChatCompletionMessageParam],
    call_type: str = "",
    trace_id: str = "",
    max_tokens: int = 16384,
    user_id: str | None = None,
    python_version: str | None = None,
    context: dict[str, Any] | None = None,
) -> LLMResponse:
    """Call LLM with OpenAI or Anthropic client."""
    from aiservice.observability.database import record_llm_call  # noqa: PLC0415

    # Create a fresh client for each request to avoid event loop issues with Django dev server
    client = get_llm_client(llm.model_type)
    if client is None:
        msg = f"LLM client for model type '{llm.model_type}' is not available"
        raise ValueError(msg)

    start_time = time.time()
    error: Exception | None = None
    result: LLMResponse | None = None

    try:
        if llm.model_type == "anthropic":
            assert isinstance(client, AsyncAnthropicFoundry)
            system_prompt_content = next((m["content"] for m in messages if m["role"] == "system"), None)
            anthropic_messages = [
                {"role": m["role"], "content": m["content"]} for m in messages if m["role"] != "system"
            ]

            kwargs: dict[str, Any] = {"model": llm.name, "messages": anthropic_messages, "max_tokens": max_tokens}
            if system_prompt_content:
                kwargs["system"] = system_prompt_content

            response = await client.messages.create(**kwargs)
            content = "".join(block.text for block in response.content if hasattr(block, "text"))

            result = LLMResponse(
                content=content,
                usage=LLMUsage(input_tokens=response.usage.input_tokens, output_tokens=response.usage.output_tokens),
                raw_response=response,
            )
        else:
            # Azure OpenAI
            assert isinstance(client, AsyncAzureOpenAI)
            # gpt-5-mini only accepts max_completion_tokens, older models use max_tokens
            if llm.name == "gpt-5-mini":
                response = await client.chat.completions.create(
                    model=llm.name, messages=messages, max_completion_tokens=max_tokens
                )
            else:
                response = await client.chat.completions.create(
                    model=llm.name, messages=messages, max_tokens=max_tokens
                )

            result = LLMResponse(
                content=response.choices[0].message.content or "",
                usage=LLMUsage(
                    input_tokens=response.usage.prompt_tokens if response.usage else 0,
                    output_tokens=response.usage.completion_tokens if response.usage else 0,
                ),
                raw_response=response,
            )
        return result

    except Exception as e:
        error = e
        raise

    finally:
        latency_ms = int((time.time() - start_time) * 1000)
        try:
            await record_llm_call(
                trace_id=trace_id,
                call_type=call_type,
                model_name=llm.name,
                messages=messages,
                user_id=user_id,
                python_version=python_version,
                context=context,
                result=result,
                error=error,
                llm_cost=calculate_llm_cost(result.raw_response, llm) if result else None,
                latency_ms=latency_ms,
            )
        except Exception as e:
            logger.warning(f"Tracing: Failed to record LLM call: {e}")


# =============================================================================
# Cost Calculation
# =============================================================================


def calculate_llm_cost(response: ChatCompletion | AnthropicMessage, llm: LLM) -> float:
    """Calculate the cost of an LLM API call based on token usage."""
    if response.usage is None:
        return 0.0

    usage = response.usage

    # Extract token counts per provider
    # OpenAI: prompt_tokens is total (cached is subset), Anthropic: counts are additive
    if llm.model_type == "anthropic":
        cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
        non_cached = usage.input_tokens + (getattr(usage, "cache_creation_input_tokens", 0) or 0)  # type: ignore[union-attr]
        output = usage.output_tokens  # type: ignore[union-attr]
    else:
        details = getattr(usage, "prompt_tokens_details", None)
        cache_read = (getattr(details, "cached_tokens", 0) or 0) if details else 0
        non_cached = usage.prompt_tokens - cache_read  # type: ignore[union-attr]
        output = usage.completion_tokens  # type: ignore[union-attr]

    input_rate = llm.input_cost or 0.0
    cached_rate = llm.cached_input_cost if llm.cached_input_cost is not None else input_rate
    output_rate = llm.output_cost or 0.0

    return (non_cached * input_rate + cache_read * cached_rate + output * output_rate) / 1_000_000


# =============================================================================
# Model Selection (based on available clients)
# =============================================================================

# Prefer OpenAI: use OpenAI if available, fall back to Anthropic
OPENAI_MODEL: LLM = OpenAI_GPT_5_Mini() if _openai_client else Anthropic_Claude_Sonnet_4_5()

# Prefer Anthropic: use Anthropic (Azure Foundry) if available, fall back to OpenAI
ANTHROPIC_MODEL: LLM = Anthropic_Claude_Sonnet_4_5() if _anthropic_client else OpenAI_GPT_5_Mini()

# Haiku model for cost-effective tasks (testgen diversity)
HAIKU_MODEL: LLM = Anthropic_Claude_Haiku_4_5() if _anthropic_client else OpenAI_GPT_5_Mini()

# Model assignments
EXPLAIN_MODEL: LLM = OPENAI_MODEL
PLAN_MODEL: LLM = OPENAI_MODEL
EXECUTE_MODEL: LLM = OPENAI_MODEL
OPTIMIZE_MODEL: LLM = OPENAI_MODEL
RANKING_MODEL: LLM = OPENAI_MODEL

REFINEMENT_MODEL: LLM = ANTHROPIC_MODEL
EXPLANATIONS_MODEL: LLM = ANTHROPIC_MODEL
OPTIMIZATION_REVIEW_MODEL: LLM = ANTHROPIC_MODEL
CODE_REPAIR_MODEL: LLM = ANTHROPIC_MODEL
ADAPTIVE_OPTIMIZE_MODEL: LLM = ANTHROPIC_MODEL
