"""
LLM Client for the Glueco Gateway.

Provides a typed client for making chat completion requests
through the gateway transport.

Example:
    >>> from glueco_sdk import create_transport
    >>> from glueco_plugin_llm import llm_client
    >>> 
    >>> transport = create_transport(proxy_url="...", app_id="...")
    >>> llm = llm_client(transport)
    >>> 
    >>> response = llm.chat_completions(
    ...     provider="groq",
    ...     model="llama-3.1-8b-instant",
    ...     messages=[{"role": "user", "content": "Hello!"}],
    ... )
    >>> print(response.content)
"""

from __future__ import annotations

import json
from typing import Any, Dict, Iterator, List, Literal, Optional, Union

from glueco_sdk import GatewayTransport

from .types import (
    ChatCompletion,
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatMessage,
)


# Supported providers
Provider = Literal["openai", "groq", "gemini", "anthropic"]


class LLMClient:
    """
    Typed LLM client for the Glueco Gateway.
    
    Provides methods for chat completions with full type safety.
    Supports multiple providers through the gateway.
    
    Attributes:
        transport: Gateway transport for making requests.
    """
    
    def __init__(self, transport: GatewayTransport) -> None:
        """
        Initialize the LLM client.
        
        Args:
            transport: Gateway transport from GatewayClient.get_transport().
        """
        self._transport = transport
    
    def chat_completions(
        self,
        provider: Provider,
        model: str,
        messages: List[Union[ChatMessage, Dict[str, Any]]],
        *,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> ChatCompletion:
        """
        Create a chat completion.
        
        Args:
            provider: LLM provider (openai, groq, gemini).
            model: Model name (e.g., "gpt-4o", "llama-3.1-8b-instant").
            messages: Conversation messages.
            temperature: Sampling temperature (0.0 to 2.0).
            max_tokens: Maximum tokens to generate.
            response_format: Response format specification.
            **kwargs: Additional provider-specific parameters.
            
        Returns:
            ChatCompletion with model response.
            
        Raises:
            GatewayError: If the request fails.
            
        Example:
            >>> response = llm.chat_completions(
            ...     provider="groq",
            ...     model="llama-3.1-8b-instant",
            ...     messages=[
            ...         {"role": "system", "content": "You are helpful."},
            ...         {"role": "user", "content": "Hello!"},
            ...     ],
            ...     temperature=0.7,
            ... )
            >>> print(response.content)
        """
        # Build request
        request = ChatCompletionRequest(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            response_format=response_format,
        )
        
        payload = request.to_dict()
        payload.update(kwargs)
        
        # Make request
        resource_id = f"llm:{provider}"
        response = self._transport.request(
            resource_id=resource_id,
            action="chat.completions",
            payload=payload,
        )
        
        return ChatCompletion.from_dict(response.data)
    
    def chat_completions_stream(
        self,
        provider: Provider,
        model: str,
        messages: List[Union[ChatMessage, Dict[str, Any]]],
        *,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> Iterator[ChatCompletionChunk]:
        """
        Create a streaming chat completion.
        
        Args:
            provider: LLM provider (openai, groq, gemini).
            model: Model name.
            messages: Conversation messages.
            temperature: Sampling temperature.
            max_tokens: Maximum tokens to generate.
            **kwargs: Additional provider-specific parameters.
            
        Yields:
            ChatCompletionChunk for each streamed response chunk.
            
        Raises:
            GatewayError: If the request fails.
            
        Example:
            >>> for chunk in llm.chat_completions_stream(
            ...     provider="groq",
            ...     model="llama-3.1-8b-instant",
            ...     messages=[{"role": "user", "content": "Tell me a story"}],
            ... ):
            ...     if chunk.choices[0].delta.content:
            ...         print(chunk.choices[0].delta.content, end="")
        """
        # Build request
        request = ChatCompletionRequest(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
        )
        
        payload = request.to_dict()
        payload.update(kwargs)
        
        # Make streaming request
        resource_id = f"llm:{provider}"
        response = self._transport.request_stream(
            resource_id=resource_id,
            action="chat.completions",
            payload=payload,
        )
        
        # Parse SSE stream
        buffer = ""
        for chunk_bytes in response.stream:
            buffer += chunk_bytes.decode("utf-8")
            
            while "\n" in buffer:
                line, buffer = buffer.split("\n", 1)
                line = line.strip()
                
                if not line:
                    continue
                
                if line.startswith("data: "):
                    data = line[6:]
                    
                    if data == "[DONE]":
                        return
                    
                    try:
                        chunk_data = json.loads(data)
                        yield ChatCompletionChunk.from_dict(chunk_data)
                    except json.JSONDecodeError:
                        continue


def llm_client(transport: GatewayTransport) -> LLMClient:
    """
    Create an LLM client from a gateway transport.
    
    This is the main entry point for the LLM plugin.
    
    Args:
        transport: Gateway transport from create_transport().
        
    Returns:
        Configured LLMClient.
        
    Example:
        >>> from glueco_sdk import create_transport
        >>> from glueco_plugin_llm import llm_client
        >>> 
        >>> transport = create_transport(proxy_url="...", app_id="...")
        >>> llm = llm_client(transport)
    """
    return LLMClient(transport)
