"""Chat completions module."""

import json
from typing import List, Dict, Optional, Union, Iterator
from .types import ChatResponse, ChatChunk, ChatMessage


class ChatCompletions:
    """Handle chat completion requests."""
    
    def __init__(self, client):
        self._client = client
    
    def create(
        self,
        model: str,
        messages: List[Dict[str, str]],
        stream: bool = False,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        frequency_penalty: Optional[float] = None,
        presence_penalty: Optional[float] = None,
        stop: Optional[Union[str, List[str]]] = None,
        **kwargs
    ) -> Union[ChatResponse, Iterator[ChatChunk]]:
        """
        Create a chat completion.
        
        Args:
            model: Model to use (e.g., "gpt-4o-mini", "claude-sonnet-4-20250514", "gemini-2.0-flash")
            messages: List of message dicts with 'role' and 'content'
            stream: Whether to stream the response
            temperature: Sampling temperature (0-2)
            max_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter
            frequency_penalty: Frequency penalty (-2 to 2)
            presence_penalty: Presence penalty (-2 to 2)
            stop: Stop sequences
            **kwargs: Additional parameters
            
        Returns:
            ChatResponse object or iterator of ChatChunk objects if streaming
            
        Example:
            >>> response = client.chat.create(
            ...     model="gpt-4o-mini",
            ...     messages=[{"role": "user", "content": "Hello!"}]
            ... )
            >>> print(response.choices[0].message.content)
        """
        payload = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        
        # Add optional parameters
        if temperature is not None:
            payload["temperature"] = temperature
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens
        if top_p is not None:
            payload["top_p"] = top_p
        if frequency_penalty is not None:
            payload["frequency_penalty"] = frequency_penalty
        if presence_penalty is not None:
            payload["presence_penalty"] = presence_penalty
        if stop is not None:
            payload["stop"] = stop
        
        # Add any extra kwargs
        payload.update(kwargs)
        
        if stream:
            return self._stream(payload)
        else:
            response = self._client._post("/chat/completions", json=payload)
            return ChatResponse(**response.json())
    
    def _stream(self, payload: dict) -> Iterator[ChatChunk]:
        """Stream chat completion chunks."""
        response = self._client._session.post(
            f"{self._client.base_url}/chat/completions",
            json=payload,
            stream=True
        )
        
        for line in response.iter_lines():
            if not line:
                continue
            
            line = line.decode('utf-8')
            
            if line.startswith('data: '):
                data = line[6:]  # Remove 'data: ' prefix
                
                if data == '[DONE]':
                    break
                
                try:
                    chunk_data = json.loads(data)
                    yield ChatChunk(**chunk_data)
                except json.JSONDecodeError:
                    continue
    
    def complete(
        self,
        model: str,
        prompt: str,
        **kwargs
    ) -> str:
        """
        Simplified chat completion - just returns the text response.
        
        Args:
            model: Model to use
            prompt: User prompt
            **kwargs: Additional parameters (temperature, max_tokens, etc.)
            
        Returns:
            Response text as string
            
        Example:
            >>> text = client.chat.complete("gpt-4o-mini", "Say hello!")
            >>> print(text)
            Hello! How can I help you today?
        """
        messages = [{"role": "user", "content": prompt}]
        response = self.create(model=model, messages=messages, **kwargs)
        return response.choices[0].message.content
