"""LLM client for Dash AI Chat API."""

from typing import Any, Dict, List, Optional

import httpx
from typing_extensions import Literal, NotRequired, TypedDict

from ._cloud_env import cloud_config
from ._oauth import OAuthClient
from .exceptions import APIError, NetworkError


class CacheControl(TypedDict):
    """Cache control for prompt caching."""

    type: Literal["ephemeral"]


class ChatMessage(TypedDict):
    """OpenAI-compatible chat message."""

    role: Literal["system", "user", "assistant", "tool"]
    content: str
    tool_calls: NotRequired[List[Dict[str, Any]]]
    tool_call_id: NotRequired[str]
    name: NotRequired[str]
    cache_control: NotRequired[CacheControl]


class ToolFunction(TypedDict):
    """Function definition within a tool."""

    name: str
    description: str
    parameters: Dict[str, Any]


class ToolDefinition(TypedDict):
    """Tool definition for LLM."""

    type: Literal["function"]
    function: ToolFunction


class ChatCompletionRequest(TypedDict):
    """Request body for chat completions."""

    model: str
    messages: List[ChatMessage]
    tools: NotRequired[List[ToolDefinition]]
    tool_choice: NotRequired[str]
    max_tokens: NotRequired[int]
    temperature: NotRequired[float]
    origin_source: str


class ToolCallFunction(TypedDict):
    """Function call details."""

    name: str
    arguments: str


class ToolCall(TypedDict):
    """Tool call from LLM response."""

    id: str
    type: Literal["function"]
    function: ToolCallFunction


class ChatCompletionMessage(TypedDict):
    """Message in chat completion response."""

    role: Literal["assistant"]
    content: Optional[str]
    tool_calls: NotRequired[List[ToolCall]]


class ChatCompletionChoice(TypedDict):
    """Choice in chat completion response."""

    index: int
    message: ChatCompletionMessage
    finish_reason: Literal["stop", "tool_calls", "length"]


class ChatCompletionResponse(TypedDict):
    """Response from chat completions API."""

    id: str
    object: str
    created: int
    model: str
    choices: List[ChatCompletionChoice]


class LLMClient:
    """Client for LLM API.

    This client communicates with the LLM API using OAuth authentication.

    Usage:
        oauth_client = OAuthClient(client_id)
        async with LLMClient(oauth_client) as llm:
            response = await llm.chat_completion(messages)
    """

    def __init__(self, oauth_client: OAuthClient):
        """Initialize the LLM client.

        Args:
            oauth_client: OAuth client for authentication
        """
        self.oauth_client = oauth_client
        self.api_base_url = cloud_config.get_api_base_url().rstrip("/")
        self._client: Optional[httpx.AsyncClient] = None

    async def __aenter__(self) -> "LLMClient":
        """Async context manager entry."""
        headers = {
            "user-agent": "PlotlyStudio",
            "Content-Type": "application/json",
        }
        access_token = await self.oauth_client.get_access_token()
        if access_token:
            headers["Authorization"] = f"Bearer {access_token}"

        self._client = httpx.AsyncClient(
            timeout=120.0,
            headers=headers,
        )
        return self

    async def __aexit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[BaseException],
        exc_tb: Optional[Any],
    ) -> None:
        """Async context manager exit."""
        if self._client:
            await self._client.aclose()
            self._client = None

    async def _refresh_token_if_needed(self, response: httpx.Response) -> bool:
        """Refresh token if response indicates authentication failure.

        Args:
            response: HTTP response to check

        Returns:
            True if token was refreshed successfully, False otherwise
        """
        # Handle both 401 (Unauthorized) and 403 (Forbidden) as potential token expiry
        if response.status_code in (401, 403) and self.oauth_client and self._client:
            try:
                new_token = await self.oauth_client.refresh_access_token()
                self._client.headers["Authorization"] = f"Bearer {new_token}"
                return True
            except Exception:
                # Refresh failed - clear stale credentials so user is prompted to re-auth
                self.oauth_client.clear_credentials()
        return False

    async def chat_completion(
        self,
        messages: List[ChatMessage],
        tools: Optional[List[ToolDefinition]] = None,
        model: str = "claude-sonnet-4",
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> ChatCompletionResponse:
        """Send chat completion request to LLM API.

        Args:
            messages: List of chat messages
            tools: Optional list of tool definitions for function calling
            model: Model to use (default: "claude-sonnet-4")
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature

        Returns:
            ChatCompletionResponse with model's response

        Raises:
            RuntimeError: If client is not initialized
            APIError: If API returns an error
            NetworkError: If network request fails
        """
        if not self._client:
            raise RuntimeError("Client not initialized. Use within async context manager.")

        url = f"https://{self.api_base_url}/api/ai/chat/completions"

        request: ChatCompletionRequest = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "origin_source": "dash_chat",
        }
        if tools:
            request["tools"] = tools
            request["tool_choice"] = "auto"

        try:
            retry_count = 0
            response = None

            while retry_count <= 1:
                response = await self._client.post(url, json=request)

                if response.status_code == 200:
                    return response.json()

                token_refreshed = await self._refresh_token_if_needed(response)
                if token_refreshed:
                    retry_count += 1
                else:
                    break

            # If we get here, request failed
            error_detail = ""
            if response is not None:
                try:
                    error_detail = response.text
                except Exception:
                    error_detail = f"Status code: {response.status_code}"

            raise APIError(
                f"LLM API error: {response.status_code if response else 'unknown'}",
                status_code=response.status_code if response else 0,
                details=error_detail,
            )

        except httpx.RequestError as e:
            raise NetworkError("Failed to reach LLM API", str(e)) from e
