import json
import logging
import os
from datetime import datetime, timezone
from typing import List, Callable

from openai import OpenAI, AzureOpenAI, AsyncOpenAI, AsyncAzureOpenAI
from openai.types.chat import ChatCompletion

from superagentx.exceptions import InvalidType
from superagentx.llm.constants import (
    DEFAULT_OPENAI_EMBED, DEFAULT_BEDROCK_EMBED, DEFAULT_OLLAMA_EMBED, DEFAULT_EMBED, DEFAULT_GEMINI_EMBED
)
from superagentx.llm.litellm import LiteLLMClient

from superagentx.llm.models import ChatCompletionParams
from superagentx.llm.openai import OpenAIClient
from superagentx.llm.types.base import LLMModelConfig
from superagentx.llm.types.response import Message, Tool
from superagentx.utils.helper import iter_to_aiter, sync_to_async
from superagentx.utils.llm_config import LLMType

logger = logging.getLogger(__name__)

_retries = 5
_deepseek_base_url = 'https://api.deepseek.com'
_anthropic_base_url = "https://api.anthropic.com/v1/"
_routeway_base_url = "https://api.routeway.ai/v1"


class LLMClient:
    """Base LLM Client to create LLM client based on model llm_type.

       Model Types: openai, azure-openai, llama, gemini, mistral, bedrock, together, groq, anthropic.

       Open AI:

       llm_config = {
        "model": "gpt-4o",
        "api_type": "openai",
      }

      llm_config = {
        "model": "gpt-4o",
        "api_type": "openai",
      }

    """

    def __init__(
            self,
            *,
            llm_config: dict,
            **kwargs
    ):
        self.llm_config = llm_config
        self.llm_config_model = LLMModelConfig(**self.llm_config)
        self.async_mode = self.llm_config_model.async_mode

        match self.llm_config_model.llm_type:
            case LLMType.OPENAI_CLIENT | LLMType.ANTHROPIC_CLIENT | LLMType.DEEPSEEK | LLMType.ROUTEWAY:
                self.client = self._init_openai_cli()
            case LLMType.AZURE_OPENAI_CLIENT:
                self.client = self._init_azure_openai_cli()
            case LLMType.BEDROCK_CLIENT:
                self.client = self._init_bedrock_cli(**kwargs)
            case LLMType.GEMINI_CLIENT:
                self.client = self._init_gemini_cli(**kwargs)
            case LLMType.OLLAMA:
                self.client = self._init_ollama_cli(**kwargs)
            case _:
                self.client = self.__init_litellm_cli()

    def _init_openai_cli(self) -> OpenAIClient:
        # Set the API Key from pydantic model class or from environment variables.

        # Determine the client class based on async_mode
        client_class = AsyncOpenAI if self.llm_config_model.async_mode else OpenAI

        base_url = None
        api_key = None

        if LLMType.ANTHROPIC_CLIENT == self.llm_config_model.llm_type:
            base_url = self.llm_config_model.base_url or _anthropic_base_url
            api_key = self.llm_config_model.api_key or os.getenv("ANTHROPIC_API_KEY")

        if LLMType.DEEPSEEK == self.llm_config_model.llm_type:
            base_url = self.llm_config_model.base_url or _deepseek_base_url
            api_key = self.llm_config_model.api_key or os.getenv("DEEPSEEK_API_KEY")

        if LLMType.ROUTEWAY == self.llm_config_model.llm_type:
            base_url = self.llm_config_model.base_url or _routeway_base_url
            api_key = self.llm_config_model.api_key or os.getenv("Routeway_API_KEY")

        if LLMType.OPENAI_CLIENT == self.llm_config_model.llm_type:
            api_key = self.llm_config_model.api_key or os.getenv("OPENAI_API_KEY")

        # Initialize the client with the API key
        cli = client_class(api_key=api_key, base_url=base_url or None)

        # Set the embed model attribute
        embed_model = self.llm_config_model.embed_model

        # Assign the client to self.client
        return OpenAIClient(
            client=cli,
            model=self.llm_config_model.model,
            embed_model=embed_model or DEFAULT_OPENAI_EMBED,
            llm_type=self.llm_config_model.llm_type
        )

    def __init_litellm_cli(self):
        return LiteLLMClient(model=self.llm_config_model.model)

    def _init_gemini_cli(self, **kwargs):
        from superagentx.llm.gemini import GeminiClient

        # Set the parameters from pydantic model class or from environment variables.
        api_key = self.llm_config_model.api_key or os.getenv("GEMINI_API_KEY")

        from google import genai

        cli = genai.Client(api_key=api_key)
        # Set the embed model attribute
        embed_model = self.llm_config_model.embed_model

        return GeminiClient(
            client=cli,
            model=self.llm_config_model.model,
            embed_model=embed_model or DEFAULT_GEMINI_EMBED,
        )

    def _init_azure_openai_cli(self) -> OpenAIClient:
        # Set the parameters from pydantic model class or from environment variables.
        api_key = self.llm_config_model.api_key or os.getenv("AZURE_OPENAI_API_KEY") or os.getenv("AZURE_API_KEY")
        base_url = self.llm_config_model.base_url or os.getenv("AZURE_ENDPOINT") or os.getenv("AZURE_API_BASE")
        azure_deployment = self.llm_config_model.model or os.getenv("AZURE_DEPLOYMENT")
        api_version = self.llm_config_model.api_version or os.getenv("API_VERSION") or os.getenv("AZURE_API_VERSION")

        # Determine the client class based on async_mode
        client_class = AsyncAzureOpenAI if self.llm_config_model.async_mode else AzureOpenAI

        # Initialize the client with the gathered configuration
        cli = client_class(
            api_key=api_key,
            azure_endpoint=base_url,
            api_version=api_version
        )

        # Assign the client to self.client
        return OpenAIClient(
            client=cli,
            model=self.llm_config_model.model,
            embed_model=self.llm_config_model.embed_model
        )

    def _init_bedrock_cli(self, **kwargs):
        import boto3
        from botocore.config import Config
        from superagentx.llm.bedrock import BedrockClient

        aws_region = kwargs.get("aws_region", None) or os.getenv("AWS_REGION") or os.getenv("AWS_REGION_NAME")

        if not aws_region:
            raise ValueError("Region is required to use the Amazon Bedrock API.")

        aws_access_key = kwargs.get("aws_access_key", None) or os.getenv("AWS_ACCESS_KEY") or os.getenv("AWS_ACCESS_KEY_ID")
        aws_secret_key = kwargs.get("aws_secret_key", None) or os.getenv("AWS_SECRET_KEY") or os.getenv("AWS_SECRET_ACCESS_KEY")

        # Initialize Bedrock client, session, and runtime
        bedrock_config = Config(
            region_name=aws_region,
            signature_version="v4",
            retries={
                "max_attempts": _retries,
                "mode": "standard"
            },
        )

        # Assign Bedrock client to self.client
        aws_cli = boto3.client(
            service_name="bedrock-runtime",
            aws_access_key_id=aws_access_key,
            aws_secret_access_key=aws_secret_key,
            config=bedrock_config
        )
        # Set the embed model attribute
        embed_model = self.llm_config_model.embed_model

        return BedrockClient(
            client=aws_cli,
            model=self.llm_config_model.model,
            embed_model=DEFAULT_BEDROCK_EMBED if not embed_model else embed_model
        )

    def _init_ollama_cli(self, **kwargs):
        from ollama import AsyncClient
        from ollama import Client as OllamaCli
        from superagentx.llm.ollama import OllamaClient

        host = kwargs.get("host", None) or os.getenv("OLLAMA_HOST")

        # Async & Sync Ollama Cli
        cli = AsyncClient(host=host) if self.llm_config_model.async_mode else OllamaCli(host=host)

        embed_model = self.llm_config_model.embed_model
        return OllamaClient(
            client=cli,
            embed_model=DEFAULT_OLLAMA_EMBED if not embed_model else embed_model,
            model=self.llm_config_model.model,
            **kwargs
        )

    def chat_completion(
            self,
            *,
            chat_completion_params: ChatCompletionParams
    ) -> ChatCompletion:
        return self.client.chat_completion(
            chat_completion_params=chat_completion_params
        )

    async def achat_completion(
            self,
            *,
            chat_completion_params: ChatCompletionParams
    ) -> ChatCompletion:
        if self.async_mode:
            return await self.client.achat_completion(
                chat_completion_params=chat_completion_params
            )
        else:
            return await sync_to_async(
                self.client.chat_completion,
                chat_completion_params=chat_completion_params
            )

    async def get_tool_json(
            self,
            *,
            func: Callable
    ) -> dict:
        return await self.client.get_tool_json(func=func)

    def embed(
            self,
            *,
            text: str,
            **kwargs
    ):
        return self.client.embed(
            text,
            **kwargs
        )

    async def aembed(
            self,
            *,
            text: str,
            **kwargs
    ):
        if self.async_mode:
            return await self.client.aembed(
                text,
                **kwargs
            )
        else:
            return await sync_to_async(
                self.client.embed,
                text,
                **kwargs
            )

    async def afunc_chat_completion(
            self,
            *,
            chat_completion_params: ChatCompletionParams
    ) -> List[Message]:

        stream = bool(chat_completion_params.stream)

        # Most models don't support streaming with tool use
        if stream:
            logger.warning(
                "Streaming is not currently supported, streaming will be disabled.",
            )
            chat_completion_params.stream = False

        if self.async_mode:
            response: ChatCompletion = await self.client.achat_completion(
                chat_completion_params=chat_completion_params
            )
        else:
            response: ChatCompletion = await sync_to_async(
                self.client.chat_completion,
                chat_completion_params=chat_completion_params
            )
        # List to store multiple Message instances
        message_instances = []

        if response:
            # Iterate over each choice and create a Message instance
            async for choice in iter_to_aiter(response.choices):
                tool_calls_data = []
                if choice.message.tool_calls:
                    tool_calls_data = [
                        Tool(
                            tool_type=tool_call.type,
                            name=tool_call.function.name,
                            arguments=json.loads(tool_call.function.arguments)  # Use json.loads for safer parsing
                        ) async for tool_call in iter_to_aiter(choice.message.tool_calls)
                    ]

                # Extract token details from usage
                usage_data = response.usage

                # Add the created Message instance to the list
                # Extract nested data once to avoid repeated attribute lookups
                msg = choice.message
                usage = usage_data
                details = getattr(usage, "completion_tokens_details", None)

                message_instances.append(
                    Message(
                        role=msg.role,
                        model=response.model,
                        content=msg.content,
                        tool_calls=tool_calls_data or None,  # Concise falsy check
                        completion_tokens=usage.completion_tokens,
                        prompt_tokens=usage.prompt_tokens,
                        total_tokens=usage.total_tokens,
                        reasoning_tokens=(
                            details.reasoning_tokens
                            if details and hasattr(details, 'reasoning_tokens')
                            else 0
                        ),
                        created=datetime.fromtimestamp(response.created, tz=timezone.utc)
                    )
                )

        return message_instances

    def count_tokens(self, chat_completion_params: ChatCompletionParams):
        return self.client.count_tokens(chat_completion_params=chat_completion_params)

    async def account_tokens(self, chat_completion_params: ChatCompletionParams):
        if self.async_mode:
            return await self.client.account_tokens(chat_completion_params=chat_completion_params)
        else:
            return await sync_to_async(self.count_tokens, chat_completion_params)