from collections.abc import Awaitable, Callable
from typing import (
    Any,
    Literal,
    ParamSpec,
    Protocol,
    TypeVar,
    overload,
    runtime_checkable,
)

from anthropic import AsyncAnthropic
from google import genai
from groq import AsyncGroq
from mirascope import llm
from mirascope.integrations.langfuse import with_langfuse
from mirascope.llm import Stream
from openai import AsyncOpenAI
from pydantic import BaseModel
from sentry_sdk.ai.monitoring import ai_track
from tenacity import retry, stop_after_attempt, wait_exponential

from src.config import settings
from src.utils.types import Providers

clients: dict[Providers, AsyncAnthropic | AsyncOpenAI | genai.Client | AsyncGroq] = {}

if settings.LLM.ANTHROPIC_API_KEY:
    anthropic = AsyncAnthropic(api_key=settings.LLM.ANTHROPIC_API_KEY)
    clients["anthropic"] = anthropic

if settings.LLM.OPENAI_API_KEY:
    openai_client = AsyncOpenAI(
        api_key=settings.LLM.OPENAI_API_KEY,
    )
    clients["openai"] = openai_client

if settings.LLM.OPENAI_COMPATIBLE_BASE_URL:
    clients["custom"] = AsyncOpenAI(
        api_key=settings.LLM.OPENAI_COMPATIBLE_API_KEY,
        base_url=settings.LLM.OPENAI_COMPATIBLE_BASE_URL,
    )

if settings.LLM.GEMINI_API_KEY:
    google = genai.Client(api_key=settings.LLM.GEMINI_API_KEY)
    clients["google"] = google

if settings.LLM.GROQ_API_KEY:
    groq = AsyncGroq(api_key=settings.LLM.GROQ_API_KEY)
    clients["groq"] = groq

providers = [
    ("Dialectic", settings.DIALECTIC.PROVIDER),
    ("Summary", settings.SUMMARY.PROVIDER),
    ("Deriver", settings.DERIVER.PROVIDER),
    ("Query Generation Provider", settings.DIALECTIC.QUERY_GENERATION_PROVIDER),
]

for provider_name, provider_value in providers:
    if provider_value not in clients:
        raise ValueError(f"Missing client for {provider_name}: {provider_value}")

P = ParamSpec("P")
T = TypeVar("T", bound=BaseModel)
T_co = TypeVar("T_co", bound=BaseModel, covariant=True)
F = TypeVar("F", bound=Callable[..., Any])


# Define protocols for different return types
@runtime_checkable
class AsyncResponseModelCallable(Protocol[P, T_co]):
    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T_co: ...


@runtime_checkable
class SyncResponseModelCallable(Protocol[P, T_co]):
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T_co: ...


@runtime_checkable
class AsyncStreamCallable(Protocol[P]):
    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Stream: ...


@runtime_checkable
class SyncStreamCallable(Protocol[P]):
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Stream: ...


@runtime_checkable
class AsyncStringCallable(Protocol[P]):
    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> str: ...


@runtime_checkable
class SyncStringCallable(Protocol[P]):
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> str: ...


# Overload for stream=True with async function
@overload
def honcho_llm_call(
    *,
    provider: Providers | None = None,
    model: str | None = None,
    track_name: str | None = None,
    response_model: type[BaseModel] | None = None,
    json_mode: bool = False,
    max_tokens: int | None = None,
    thinking_budget_tokens: int | None = None,
    enable_retry: bool = True,
    retry_attempts: int = 3,
    stream: Literal[True],
    **extra_call_params: Any,
) -> Callable[[Callable[P, Awaitable[Any]]], AsyncStreamCallable[P]]: ...


# Overload for response_model with async function
@overload
def honcho_llm_call(
    *,
    provider: Providers | None = None,
    model: str | None = None,
    track_name: str | None = None,
    response_model: type[T],
    json_mode: bool = False,
    max_tokens: int | None = None,
    thinking_budget_tokens: int | None = None,
    enable_retry: bool = True,
    retry_attempts: int = 3,
    stream: Literal[False] = False,
    **extra_call_params: Any,
) -> Callable[[Callable[P, Awaitable[Any]]], AsyncResponseModelCallable[P, T]]: ...


# Overload for no response_model with async function
@overload
def honcho_llm_call(
    *,
    provider: Providers | None = None,
    model: str | None = None,
    track_name: str | None = None,
    response_model: None = None,
    json_mode: bool = False,
    max_tokens: int | None = None,
    thinking_budget_tokens: int | None = None,
    enable_retry: bool = True,
    retry_attempts: int = 3,
    stream: Literal[False] = False,
    **extra_call_params: Any,
) -> Callable[[Callable[P, Awaitable[Any]]], AsyncStringCallable[P]]: ...


# Generic overload for sync functions (fallback)
@overload
def honcho_llm_call(
    *,
    provider: Providers | None = None,
    model: str | None = None,
    track_name: str | None = None,
    response_model: type[BaseModel] | None = None,
    json_mode: bool = False,
    max_tokens: int | None = None,
    thinking_budget_tokens: int | None = None,
    enable_retry: bool = True,
    retry_attempts: int = 3,
    stream: bool = False,
    **extra_call_params: Any,
) -> Callable[[Callable[P, Any]], Callable[P, Any]]: ...


def honcho_llm_call(
    provider: Providers | None = None,
    model: str | None = None,
    track_name: str | None = None,
    response_model: type[BaseModel] | None = None,
    json_mode: bool = False,
    max_tokens: int | None = None,
    thinking_budget_tokens: int | None = None,
    enable_retry: bool = True,
    retry_attempts: int = 3,
    stream: bool = False,
    **extra_call_params: Any,
) -> Any:
    """
    Consolidated decorator for LLM calls that handles provider-specific configurations.

    This decorator automatically:
    - Handles both sync and async functions seamlessly
    - Applies retry logic with exponential backoff
    - Adds AI tracking for Sentry
    - Integrates with Langfuse for observability
    - Builds provider-specific call parameters
    - Handles client selection from the global clients dict

    Args:
        provider: The LLM provider to use (e.g., "anthropic", "google", "openai")
        model: The model to use
        track_name: Name for AI tracking (e.g., "Critical Analysis Call")
        response_model: Optional Pydantic model for structured responses
        json_mode: Whether to enable JSON mode (for providers that support it)
        max_tokens: Maximum tokens for the response
        thinking_budget_tokens: Budget for thinking tokens (Anthropic only)
        enable_retry: Whether to enable retry logic (default: True)
        retry_attempts: Number of retry attempts (default: 3)
        stream: Whether to enable streaming responses (default: False)
        **extra_call_params: Additional provider-specific parameters

    Returns:
        A decorator that returns:
        - For async functions: Callable[P, Awaitable[T]] where T is Stream, response_model, or str
        - For sync functions: Callable[P, T] where T is Stream, response_model, or str

    Note: Type annotations may be needed at the call site for proper type checking.

    Example (async function):
        @honcho_llm_call(
            provider=settings.DERIVER.PROVIDER,
            model=settings.DERIVER.MODEL,
            track_name="Critical Analysis Call",
            response_model=ReasoningResponse,
            json_mode=True,
            max_tokens=settings.DERIVER.MAX_OUTPUT_TOKENS,
        )
        async def analyze(context: str, query: str):
            return prompt_template(context, query)

    Example (sync function):
        @honcho_llm_call(
            provider="openai",
            model="gpt-4",
            max_tokens=1000,
        )
        def generate_summary(text: str) -> str:
            return f"Summarize: {text}"

        # Call synchronously
        result = generate_summary("Long text here...")
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        # Handle special case for custom provider
        # Custom providers use OpenAI-compatible endpoints, so we resolve to "openai" for the provider name
        # but keep the original "custom" for client lookup
        resolved_provider = "openai" if provider == "custom" else provider

        # Build provider-specific call params
        call_params: dict[str, Any] = {}

        if resolved_provider == "google":
            # Google uses 'config' parameter
            config: dict[str, Any] = {}
            if max_tokens:
                config["max_output_tokens"] = max_tokens
            if json_mode or response_model:
                config["response_mime_type"] = "application/json"
                if response_model:
                    config["response_schema"] = response_model
            if config:
                call_params["config"] = config
        elif resolved_provider == "anthropic":
            # Anthropic uses thinking params and max_tokens
            if thinking_budget_tokens:
                call_params["thinking"] = {
                    "type": "enabled",
                    "budget_tokens": thinking_budget_tokens,
                }
            if max_tokens:
                call_params["max_tokens"] = max_tokens
        else:
            # Other providers just use max_tokens
            if max_tokens:
                call_params["max_tokens"] = max_tokens

        # Merge with any extra call params
        call_params.update(extra_call_params)

        # Build kwargs for llm.call
        llm_kwargs: dict[str, Any] = {}
        if resolved_provider and provider:
            llm_kwargs["provider"] = resolved_provider
            llm_kwargs["client"] = clients[
                provider
            ]  # Use original provider for client lookup
        if model:
            llm_kwargs["model"] = model
        if response_model:
            llm_kwargs["response_model"] = response_model
        if json_mode:
            llm_kwargs["json_mode"] = json_mode
        if stream:
            llm_kwargs["stream"] = stream
        if call_params:
            llm_kwargs["call_params"] = call_params

        # Apply decorators in order
        decorated: Any = func

        # Apply llm.call
        decorated = llm.call(**llm_kwargs)(decorated)  # pyright: ignore

        # Apply langfuse
        decorated = with_langfuse()(decorated)  # pyright: ignore

        # Apply AI tracking if name provided
        if track_name:
            decorated = ai_track(track_name)(decorated)

        # Apply retry logic if enabled
        if enable_retry:
            decorated = retry(
                stop=stop_after_attempt(retry_attempts),
                wait=wait_exponential(multiplier=1, min=4, max=10),
            )(decorated)

        return decorated

    return decorator