from typing import List, Callable, Type, Union, Optional, Any, Dict, AsyncIterator
import os
import anthropic
import openai
from google import genai
from google.genai import types as genai_types
import json
import backoff
from typing import Dict, Any
from uuid import uuid4
from pydantic import BaseModel
from cache_on_disk import DCache
from .dtypes import (
    ChatMessage,
    MessageRole,
    anthropic_format,
    openai_format,
    genai_format,
    TextBlock,
    ToolDefinition,
    ReasoningConfig,
    ThinkingBlock,
)

from dotenv import load_dotenv

load_dotenv()


class Provider:
    def __init__(self, get_response: Callable[..., Any], models: List[str]) -> None:
        self.models = models
        self.get_response = get_response


# ---------------------------------------------------------------------------
# Anthropic provider
# ---------------------------------------------------------------------------

anthr = anthropic.AsyncAnthropic()


async def get_response_anthropic(
    messages: List[ChatMessage],
    tools: Optional[List[ToolDefinition]],
    response_format: Optional[Dict[str, Any]] = None,
    reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
    **kwargs: Any,
) -> ChatMessage:
    if response_format is not None:
        raise NotImplementedError(
            "Structured output is not supported for Anthropic models"
        )

    kwargs = anthropic_format(messages, tools, reasoning=reasoning, **kwargs)
    kwargs["timeout"] = 599
    resp = await anthr.messages.create(**kwargs)

    return ChatMessage.from_anthropic(resp.content)


# ---------------------------------------------------------------------------
# OpenAI provider factory
# ---------------------------------------------------------------------------


def get_response_factory(oai: openai.AsyncOpenAI) -> Callable[..., Any]:
    async def get_response_openai(
        messages: List[ChatMessage],
        tools: Optional[List[ToolDefinition]],
        response_format: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
        reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
        **kwargs: Any,
    ) -> ChatMessage:
        kwargs = openai_format(messages, tools, reasoning=reasoning, **kwargs)

        if "model" not in kwargs:
            raise ValueError("'model' is required for OpenAI completions")

        # Handle structured output
        if (
            response_format is not None
            and isinstance(response_format, type)
            and issubclass(response_format, BaseModel)
        ):
            resp = await oai.chat.completions.parse(
                response_format=response_format, **kwargs
            )
            response = ChatMessage.from_openai(resp.choices[0].message)
            response.parsed = resp.choices[0].message.parsed
            return response

        # Regular completion
        if response_format is not None:
            kwargs["response_format"] = response_format

        resp = await oai.chat.completions.create(**kwargs)
        return ChatMessage.from_openai(resp.choices[0].message)

    return get_response_openai


# ---------------------------------------------------------------------------
# Google GenAI provider
# ---------------------------------------------------------------------------


async def get_response_genai(
    messages: List[ChatMessage],
    tools: Optional[List[ToolDefinition]],
    response_format: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
    reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
    **kwargs: Any,
) -> ChatMessage:

    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        raise ValueError(
            "GEMINI_API_KEY or GOOGLE_API_KEY environment variable is required"
        )

    client = genai.Client(api_key=api_key)
    request_kwargs = genai_format(messages, tools, reasoning=reasoning)

    # Build config
    config_params: Dict[str, Any] = {}
    for k, v in kwargs.items():
        if k not in [
            "contents",
            "tools",
            "system_instruction",
            "model",
            "thinking_budget",
        ]:
            if k == "max_tokens":
                config_params["max_output_tokens"] = v
            elif k in ["temperature", "top_p"]:
                config_params[k] = v

    # Handle structured output
    if (
        response_format is not None
        and isinstance(response_format, type)
        and issubclass(response_format, BaseModel)
    ):
        config_params.update(
            {
                "response_mime_type": "application/json",
                "response_schema": response_format,
            }
        )

    # Add tools and system instruction from request
    if "tools" in request_kwargs and request_kwargs["tools"]:
        config_params["tools"] = request_kwargs["tools"]
    if "system_instruction" in request_kwargs and request_kwargs["system_instruction"]:
        config_params["system_instruction"] = request_kwargs["system_instruction"]

    # Handle thinking budget
    if "thinking_budget" in request_kwargs:
        config_params["thinking_config"] = genai_types.ThinkingConfig(
            thinking_budget=request_kwargs["thinking_budget"],
            include_thoughts=True,  # Include thought summaries
        )

    config = (
        genai_types.GenerateContentConfig(**config_params) if config_params else None
    )
    # Make request
    response = await client.aio.models.generate_content(
        model=kwargs.get("model", "gemini-2.5-pro"),
        contents=request_kwargs["contents"],
        config=config,
    )
    # Convert response
    chat_response = ChatMessage.from_genai(response)

    # Handle structured output parsing
    if (
        response_format is not None
        and isinstance(response_format, type)
        and issubclass(response_format, BaseModel)
    ):
        if hasattr(response, "parsed") and response.parsed:
            chat_response.parsed = response.parsed
        elif response.text:
            try:
                parsed_data = json.loads(response.text)
                chat_response.parsed = response_format(**parsed_data)
            except Exception:
                pass

    return chat_response


# ---------------------------------------------------------------------------
# Provider registration
# ---------------------------------------------------------------------------

providers: List[Provider] = []

# Anthropic
try:
    _available_anthropic_models = [
        m.id for m in anthropic.Anthropic().models.list(limit=1000).data
    ]
    providers.append(
        Provider(get_response_anthropic, models=_available_anthropic_models)
    )
except Exception:
    pass

# OpenRouter
if "OPENROUTER_API_KEY" in os.environ:
    providers.append(
        Provider(
            get_response_factory(
                openai.AsyncOpenAI(
                    api_key=os.environ["OPENROUTER_API_KEY"],
                    base_url="https://openrouter.ai/api/v1",
                )
            ),
            models=["qwen/qwen3-235b-a22b"],
        )
    )

# OpenAI
if "OPENAI_API_KEY" in os.environ:
    try:
        _available_openai_models = [
            m.id
            for m in openai.OpenAI().models.list().data
            if m.id.startswith("gpt") or m.id.startswith("o")
        ]
        providers.append(
            Provider(
                get_response_factory(openai.AsyncOpenAI()),
                models=_available_openai_models,
            )
        )
    except Exception:
        pass

# Google GenAI
if "GEMINI_API_KEY" in os.environ or "GOOGLE_API_KEY" in os.environ:
    providers.append(
        Provider(
            get_response_genai,
            models=["gemini-2.5-pro", "gemini-2.5-flash"],
        )
    )


# ---------------------------------------------------------------------------
# Main API
# ---------------------------------------------------------------------------
async def get_response(
    model: str,
    messages: List[ChatMessage],
    tools: Optional[List[ToolDefinition]] = None,
    response_format: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
    reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
    **kwargs: Any,
) -> ChatMessage:
    # Convert dict to ReasoningConfig if needed (centralized conversion)
    if reasoning and isinstance(reasoning, dict):
        reasoning = ReasoningConfig(**reasoning)

    all_models = []
    for provider in providers:
        all_models.extend(provider.models)
        if model in provider.models:
            response = await provider.get_response(
                model=model,
                messages=messages,
                tools=tools,
                response_format=response_format,
                reasoning=reasoning,
                **kwargs,
            )

            # Handle empty responses
            if len(response.content) == 0:
                messages = messages + [
                    ChatMessage(
                        role=MessageRole.user,
                        content=[TextBlock(text="Please continue.")],
                    )
                ]
                response = await provider.get_response(
                    model=model,
                    messages=messages,
                    tools=tools,
                    response_format=response_format,
                    reasoning=reasoning,
                    **kwargs,
                )

            assert len(response.content) > 0, "Response content is empty"
            return response

    raise ValueError(
        f"Model '{model}' not supported by any provider. Supported models: {all_models}"
    )


@backoff.on_exception(
    wait_gen=backoff.expo,
    exception=(
        openai.RateLimitError,
        openai.APIConnectionError,
        openai.APITimeoutError,
        openai.InternalServerError,
        anthropic.APIConnectionError,
        anthropic.RateLimitError,
        anthropic.APIStatusError,
        genai.errors.ClientError,
        genai.errors.ServerError,
        TypeError,
        AssertionError,
    ),
    max_value=60,
    factor=1.5,
    on_backoff=lambda details: print(f"Retrying... {details['exception']}"),
)
async def get_response_with_backoff(
    model: str,
    messages: List[ChatMessage],
    tools: Optional[List[ToolDefinition]] = None,
    response_format: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
    reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
    **kwargs: Any,
) -> ChatMessage:
    """Get a response from an LLM with exponential backoff retry logic for various API errors.

    Args:
        model (str): The name of the model to use
        messages (List[ChatMessage]): The conversation history
        tools (Optional[List[ToolDefinition]]): Optional list of tools/functions the model can use
        response_format (Optional[Dict]): Optional response format specification
        reasoning (Optional[Union[ReasoningConfig, Dict]]): Optional reasoning/thinking configuration
        **kwargs: Additional keyword arguments passed to the underlying API

    Returns:
        ChatResponse: The model's response

    Raises:
        ValueError: If the specified model is not supported by any provider
    """
    # Convert dict to ReasoningConfig if needed
    if reasoning and isinstance(reasoning, dict):
        reasoning = ReasoningConfig(**reasoning)

    return await get_response(
        model,
        messages,
        tools=tools,
        response_format=response_format,
        reasoning=reasoning,
        **kwargs,
    )


dcache = DCache(cache_dir=os.path.expanduser("~/.cache/localrouter"))


@dcache(required_kwargs=["cache_seed"])
async def get_response_cached(
    model: str,
    messages: List[ChatMessage],
    tools: Optional[List[ToolDefinition]] = None,
    response_format: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
    reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
    cache_seed: Optional[int] = None,
    **kwargs: Any,
) -> ChatMessage:
    """Get a response from an LLM. Cache results on disk if called with same arguments and seed.

    Args:
        model (str): The name of the model to use
        messages (List[ChatMessage]): The conversation history
        tools (Optional[List[ToolDefinition]]): Optional list of tools/functions the model can use
        response_format (Optional[Dict]): Optional response format specification
        reasoning (Optional[Union[ReasoningConfig, Dict]]): Optional reasoning/thinking configuration
        cache_seed (int): if set, use cached responses
        **kwargs: Additional keyword arguments passed to the underlying API

    Returns:
        ChatResponse: The model's response, either from cache or freshly generated

    Raises:
        ValueError: If the specified model is not supported by any provider
    """
    # Convert dict to ReasoningConfig if needed
    if reasoning and isinstance(reasoning, dict):
        reasoning = ReasoningConfig(**reasoning)

    return await get_response(
        model,
        messages,
        tools=tools,
        response_format=response_format,
        reasoning=reasoning,
        **kwargs,
    )


@dcache(required_kwargs=["cache_seed"])
async def get_response_cached_with_backoff(
    model: str,
    messages: List[ChatMessage],
    tools: Optional[List[ToolDefinition]] = None,
    response_format: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
    reasoning: Optional[Union[ReasoningConfig, Dict[str, Any]]] = None,
    cache_seed: Optional[int] = None,
    **kwargs: Any,
) -> ChatMessage:
    """Get a response from an LLM. Cache results on disk if called with same arguments and seed. When no cached result is found, use backoff.

    Args:
        model (str): The name of the model to use
        messages (List[ChatMessage]): The conversation history
        tools (Optional[List[ToolDefinition]]): Optional list of tools/functions the model can use
        response_format (Optional[Dict]): Optional response format specification
        reasoning (Optional[Union[ReasoningConfig, Dict]]): Optional reasoning/thinking configuration
        cache_seed (int): if set, use cached responses
        **kwargs: Additional keyword arguments passed to the underlying API

    Returns:
        ChatResponse: The model's response, either from cache or freshly generated

    Raises:
        ValueError: If the specified model is not supported by any provider
    """
    # Convert dict to ReasoningConfig if needed
    if reasoning and isinstance(reasoning, dict):
        reasoning = ReasoningConfig(**reasoning)

    return await get_response_with_backoff(
        model,
        messages,
        tools=tools,
        response_format=response_format,
        reasoning=reasoning,
        **kwargs,
    )
