"""LLM module for OpenAI-compatible API calls."""

import json
from typing import Any, Literal, Optional
import httpx
from environs import Env
from rich.console import Console
from rich.prompt import Prompt

from yier_agents.src.message import Message, ToolCall, LLMResponse

# Initialize rich console for colored output
console = Console()


class LLM:
    """LLM class for making OpenAI-compatible API calls."""

    def __init__(
        self,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        timeout: float = 300.0
    ):
        """Initialize LLM client.

        Args:
            base_url: Base URL for the API endpoint (defaults to YIER_BASE_URL env var)
            api_key: API key for authentication (defaults to YIER_API_KEY env var)
            model: Model name to use (defaults to YIER_MODEL env var)
            timeout: Request timeout in seconds
        """
        env = Env()
        env.read_env()

        # Get base_url from env or prompt if not provided
        if base_url is None:
            base_url = env.str("YIER_BASE_URL", None)
            if base_url is None:
                base_url = Prompt.ask("[bold cyan]🌐 Enter base URL[/bold cyan]")

        # Get api_key from env or prompt if not provided
        if api_key is None:
            api_key = env.str("YIER_API_KEY", None)
            if api_key is None:
                api_key = Prompt.ask("[bold yellow]🔑 Enter API key[/bold yellow]", password=True)

        # Get model from env or prompt if not provided
        if model is None:
            model = env.str("YIER_MODEL", None)
            if model is None:
                model = Prompt.ask("[bold green]🤖 Enter model name[/bold green]")

        self.base_url = base_url
        self.api_key = api_key
        self.model = model
        self.timeout = timeout

    async def complete(
        self,
        messages: list[Message],
        tools: Optional[list[dict[str, Any]]] = None,
        max_tokens: int = 4096,
        temperature: float = 0.7
    ) -> LLMResponse:
        """Make a completion request to the LLM API.

        Args:
            messages: List of messages in the conversation
            tools: Optional list of tool definitions
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature

        Returns:
            LLMResponse object containing the response
        """
        # Construct request payload
        payload = {
            "model": self.model,
            "messages": [self._message_to_dict(msg) for msg in messages],
            "max_tokens": max_tokens,
            "temperature": temperature
        }

        if tools:
            payload["tools"] = tools

        # Make API request
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.post(
                f"{self.base_url}/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json"
                },
                json=payload
            )

        # Parse response
        response_data = response.json()
        choice = response_data["choices"][0]
        api_message = choice["message"]

        # Extract content, reasoning content and tool calls
        content = api_message.get("content")
        reasoning_content = api_message.get("reasoning")  # 支持思考过程
        raw_tool_calls = api_message.get("tool_calls", [])

        # Parse tool calls into ToolCall objects
        tool_calls = None
        if raw_tool_calls:
            tool_calls = []
            for tc in raw_tool_calls:
                tool_calls.append(
                    ToolCall(
                        id=tc["id"],
                        name=tc["function"]["name"],
                        arguments=json.loads(tc["function"]["arguments"])
                    )
                )

        # Normalize finish reason
        raw_finish_reason = choice.get("finish_reason")
        finish_reason = self._normalize_finish_reason(
            raw_finish_reason,
            has_tool_calls=tool_calls is not None
        )

        # Create Message object
        message = Message(
            role="assistant",
            content=content,
            reasoning_content=reasoning_content,
            tool_calls=tool_calls
        )

        # Extract usage if available (total tokens only)
        usage_data = response_data.get("usage")
        usage = usage_data.get("total_tokens") if usage_data else None

        return LLMResponse(
            message=message,
            finish_reason=finish_reason,
            usage=usage
        )

    def _normalize_finish_reason(
        self,
        raw_reason: Optional[str],
        has_tool_calls: bool
    ) -> Literal["stop", "tool_calls", "length"]:
        """Normalize finish reason to standard values.

        Args:
            raw_reason: Raw finish reason from API
            has_tool_calls: Whether the response has tool calls

        Returns:
            Normalized finish reason
        """
        if has_tool_calls:
            return "tool_calls"
        if raw_reason == "length":
            return "length"
        return "stop"

    def _message_to_dict(self, msg: Message) -> dict[str, Any]:
        """Convert Message object to OpenAI API format.

        Args:
            msg: Message object to convert

        Returns:
            Dictionary in OpenAI API format
        """
        result = {
            "role": msg.role,
            "content": msg.content
        }

        # Add reasoning content if present (for models that support it)
        if msg.reasoning_content:
            result["reasoning_content"] = msg.reasoning_content

        # Add tool calls if present (assistant message with tool calls)
        if msg.tool_calls:
            tool_calls_list = []
            for tc in msg.tool_calls:
                tool_calls_list.append({
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.name,
                        "arguments": json.dumps(tc.arguments)
                    }
                })
            result["tool_calls"] = tool_calls_list

        # Add tool_call_id if present (tool result message)
        if msg.tool_call_id:
            result["tool_call_id"] = msg.tool_call_id

        return result
