import logging
import time
import typing as T
from enum import Enum

import openai
from pydantic import BaseModel, Field

from llm import LLM, LLMConfig
from llm.function_schemas import convert_function_calls_to_actions
from llm.output_model import CortexOutputModel
from prometheus import om1_llm_latency, om1_llm_latency_last
from providers.avatar_llm_state_provider import AvatarLLMState
from providers.httpx import get_async_httpx_client
from providers.llm_history_manager import LLMHistoryManager

R = T.TypeVar("R", bound=BaseModel)


class XAIModel(str, Enum):
    """Available XAI models."""

    GROK_2_LATEST = "grok-2-latest"
    GROK_3_BETA = "grok-3-beta"
    GROK_4_LATEST = "grok-4-latest"
    GROK_4 = "grok-4"


class XAIConfig(LLMConfig):
    """XAI-specific configuration with model enum."""

    base_url: T.Optional[str] = Field(
        default="https://api.openmind.com/api/core/xai",
        description="Base URL for the XAI API endpoint",
    )
    model: T.Optional[T.Union[XAIModel, str]] = Field(
        default=XAIModel.GROK_4_LATEST,
        description="XAI model to use",
    )


class XAILLM(LLM[R]):
    """
    XAI LLM implementation using OpenAI-compatible API.

    Handles authentication and response parsing for XAI endpoints.
    """

    def __init__(
        self,
        config: XAIConfig,
        available_actions: T.Optional[T.List] = None,
    ):
        """
        Initialize the XAI LLM instance.

        Parameters
        ----------
        config : XAIConfig
            Configuration settings for the LLM.
        available_actions : list[AgentAction], optional
            List of available actions for function calling.
        """
        super().__init__(config, available_actions)

        if not config.api_key:
            raise ValueError("config file missing api_key")
        if not config.model:
            self._config.model = "grok-4-latest"

        self.base_url = config.base_url or "https://api.openmind.com/api/core/xai"
        self._client = openai.AsyncOpenAI(
            base_url=self.base_url,
            api_key=config.api_key,
            http_client=get_async_httpx_client(),
        )

        # Initialize history manager
        self.history_manager = LLMHistoryManager(self._config, self._client)

    @AvatarLLMState.trigger_thinking()
    @LLMHistoryManager.update_history()
    async def ask(self, prompt: str, messages: T.Optional[T.List[T.Dict[str, str]]] = None) -> T.Optional[R]:
        """
        Execute LLM query and parse response.

        Parameters
        ----------
        prompt : str
            The input prompt to send to the model.
        messages : List[Dict[str, str]], optional
            List of message dictionaries to send to the model.

        Returns
        -------
        R or None
            Parsed response matching the output_model structure, or None if
            parsing fails.
        """
        if messages is None:
            messages = []
        try:
            logging.debug(f"XAI LLM input: {prompt}")
            logging.debug(f"XAI LLM messages: {messages}")

            llm_start_time = time.time()
            self.io_provider.set_llm_prompt(prompt)

            formatted_messages = [
                {"role": msg.get("role", "user"), "content": msg.get("content", "")} for msg in messages
            ]
            formatted_messages.append({"role": "user", "content": prompt})

            response = await self._client.chat.completions.create(
                model=self._config.model or XAIModel.GROK_4,
                messages=T.cast(T.Any, formatted_messages),
                tools=T.cast(T.Any, self.function_schemas),
                tool_choice="auto",
                timeout=self._config.timeout,
            )

            if not response.choices:
                logging.warning("xAI API returned empty choices")
                return None

            message = response.choices[0].message
            latency = time.time() - llm_start_time
            om1_llm_latency.labels(
                model=str(self._config.model or XAIModel.GROK_4), endpoint=str(self.base_url)
            ).observe(latency)
            om1_llm_latency_last.labels(
                model=str(self._config.model or XAIModel.GROK_4), endpoint=str(self.base_url)
            ).set(latency)

            if message.tool_calls:
                logging.info(f"Received {len(message.tool_calls)} function calls")
                logging.info(f"Function calls: {message.tool_calls}")

                function_call_data = [
                    {
                        "function": {
                            "name": getattr(tc, "function").name,
                            "arguments": getattr(tc, "function").arguments,
                        }
                    }
                    for tc in message.tool_calls
                ]

                actions = convert_function_calls_to_actions(function_call_data)

                result = CortexOutputModel(actions=actions)
                self.tracer.gauge(
                    llm_input=prompt,
                    llm_output=[{"type": a.type, "value": a.value} for a in actions],
                )
                logging.info(f"XAI LLM function call output: {result}")
                return T.cast(R, result)

            self.tracer.gauge(llm_input=prompt, llm_output=[])
            return None
        except Exception as e:
            logging.error(f"XAI API error: {e}")
            return None
