"""
Main client implementation for the Inferno API.
"""

import json
import requests
from typing import Dict, Any, List, Optional, Union, Iterator, overload, Literal, cast
import time
from dataclasses import asdict, is_dataclass

from .config import InfernoConfig
from .exceptions import InfernoAPIError, InfernoConnectionError, InfernoTimeoutError
from .models import Tool, ToolChoice # Import necessary models
from .utils import (
    MessagesInput, # Import the type alias
    prepare_stop_sequences
)


class InfernoClient:
    """
    Main client for interacting with the Inferno API.
    
    This client is designed to be compatible with the OpenAI Python client,
    allowing for easy migration from OpenAI to Inferno.
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        api_version: Optional[str] = None,
        organization: Optional[str] = None,
        timeout: float = 60.0,
        max_retries: int = 3,
        default_headers: Optional[Dict[str, str]] = None,
    ):
        """
        Initialize the Inferno client.
        
        Args:
            api_key: API key for authentication (not used by Inferno but kept for OpenAI compatibility)
            api_base: Base URL for the Inferno API
            api_version: API version (not used by Inferno but kept for OpenAI compatibility)
            organization: Organization ID (not used by Inferno but kept for OpenAI compatibility)
            timeout: Request timeout in seconds
            max_retries: Maximum number of retries for failed requests
            default_headers: Default headers to include in all requests
        """
        self.config = InfernoConfig(
            api_key=api_key,
            api_base=api_base,
            api_version=api_version,
            organization=organization,
            timeout=timeout,
            max_retries=max_retries,
            default_headers=default_headers,
        )
        
        # Initialize API resources
        self.completions = Completion(self)
        self.chat = ChatCompletion(self)
        self.embeddings = Embedding(self)
        self.models = Model(self)
    
    @overload
    def request(
        self,
        method: str,
        path: str,
        params: Optional[Dict[str, Any]] = None,
        json_data: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        stream: Literal[True] = True,
        timeout: Optional[float] = None,
    ) -> Iterator[Dict[str, Any]]: ...

    @overload
    def request(
        self,
        method: str,
        path: str,
        params: Optional[Dict[str, Any]] = None,
        json_data: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        stream: Literal[False] = False,
        timeout: Optional[float] = None,
    ) -> Dict[str, Any]: ...

    def request(
        self,
        method: str,
        path: str,
        params: Optional[Dict[str, Any]] = None,
        json_data: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        stream: bool = False,
        timeout: Optional[float] = None,
    ) -> Union[Dict[str, Any], Iterator[Dict[str, Any]]]:
        """
        Make a request to the Inferno API.
        
        Args:
            method: HTTP method (GET, POST, etc.)
            path: API path
            params: Query parameters
            json_data: JSON data for the request body
            headers: Additional headers
            stream: Whether to stream the response
            timeout: Request timeout in seconds
            
        Returns:
            Union[Dict[str, Any], Iterator[Dict[str, Any]]]: API response
            
        Raises:
            InfernoAPIError: If the API returns an error
            InfernoConnectionError: If there's a connection error
            InfernoTimeoutError: If the request times out
        """
        # Build the request URL
        url = f"{self.config.api_base.rstrip('/')}/{path.lstrip('/')}"
        
        # Prepare headers
        request_headers = self.config.get_headers(headers)
        
        # Set timeout
        request_timeout = timeout if timeout is not None else self.config.timeout
        
        # Make the request with retries
        retries = 0
        while True:
            try:
                response = requests.request(
                    method=method,
                    url=url,
                    params=params,
                    json=json_data,
                    headers=request_headers,
                    timeout=request_timeout,
                    stream=stream,
                )
                
                # Handle streaming response
                if stream and response.status_code == 200:
                    return self._handle_streaming_response(response)
                
                # Handle non-streaming response
                if response.status_code != 200:
                    error_message = f"API request failed with status code: {response.status_code}"
                    try:
                        error_data = response.json()
                        if "error" in error_data:
                            error_message = error_data["error"].get("message", error_message)
                    except Exception:
                        pass
                    
                    raise InfernoAPIError(response.status_code, error_message, response)
                
                return response.json()
                
            except requests.exceptions.Timeout:
                if retries >= self.config.max_retries:
                    raise InfernoTimeoutError(request_timeout)
                
                retries += 1
                time.sleep(min(2 ** retries, 30))  # Exponential backoff
                
            except requests.exceptions.RequestException as e:
                if retries >= self.config.max_retries:
                    raise InfernoConnectionError(str(e), e)
                
                retries += 1
                time.sleep(min(2 ** retries, 30))  # Exponential backoff
    
    def _handle_streaming_response(self, response: requests.Response) -> Iterator[Dict[str, Any]]:
        """
        Handle a streaming response from the API.
        
        Args:
            response: Streaming response
            
        Returns:
            Iterator[Dict[str, Any]]: Stream of parsed responses
            
        Raises:
            InfernoAPIError: If the API returns an error
        """
        for line in response.iter_lines():
            if not line:
                continue
            
            line_text = line.decode("utf-8")
            
            # Check for data prefix and parse JSON
            if line_text.startswith("data: "):
                line_text = line_text[6:]
                
                # Check for end of stream
                if line_text.strip() == "[DONE]":
                    break
                
                try:
                    data = json.loads(line_text)
                    yield data
                except json.JSONDecodeError:
                    continue


class Completion:
    """
    Completions API resource.
    """
    
    def __init__(self, client: InfernoClient):
        """
        Initialize the Completions API resource.
        
        Args:
            client: Inferno client
        """
        self.client = client
    
    def create(
        self,
        model: str,
        prompt: Union[str, List[str]],
        suffix: Optional[str] = None,
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.95,
        n: int = 1,
        stream: bool = False,
        logprobs: Optional[int] = None,
        echo: bool = False,
        stop: Optional[Union[str, List[str]]] = None,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0,
        best_of: int = 1,
        logit_bias: Optional[Dict[str, float]] = None,
        user: Optional[str] = None,
    ) -> Union[Dict[str, Any], Iterator[Dict[str, Any]]]:
        """
        Create a completion.
        
        Args:
            model: ID of the model to use
            prompt: Prompt to generate completions for
            suffix: Suffix that comes after a completion of inserted text
            max_tokens: Maximum number of tokens to generate
            temperature: Sampling temperature
            top_p: Nucleus sampling parameter
            n: Number of completions to generate
            stream: Whether to stream the response
            logprobs: Include log probabilities of the top n tokens
            echo: Echo back the prompt in addition to the completion
            stop: Up to 4 sequences where the API will stop generating further tokens
            presence_penalty: Penalty for token presence
            frequency_penalty: Penalty for token frequency
            best_of: Generate best_of completions server-side and return the "best"
            logit_bias: Modify the likelihood of specified tokens appearing in the completion
            user: A unique identifier representing your end-user
            
        Returns:
            Union[Dict[str, Any], Iterator[Dict[str, Any]]]: Completion response
        """
        # Prepare request data
        json_data = {
            "model": model,
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "stream": stream,
        }
        
        # Add optional parameters
        if suffix is not None:
            json_data["suffix"] = suffix
        
        if n != 1:
            json_data["n"] = n
            
        if logprobs is not None:
            json_data["logprobs"] = logprobs
            
        if echo:
            json_data["echo"] = echo
            
        if stop is not None:
            json_data["stop"] = prepare_stop_sequences(stop)
            
        if presence_penalty != 0.0:
            json_data["presence_penalty"] = presence_penalty
            
        if frequency_penalty != 0.0:
            json_data["frequency_penalty"] = frequency_penalty
            
        if best_of != 1:
            json_data["best_of"] = best_of
            
        if logit_bias is not None:
            json_data["logit_bias"] = logit_bias
            
        if user is not None:
            json_data["user"] = user
        
        # Make the request
        return self.client.request(
            method="POST",
            path="/completions",
            json_data=json_data,
            stream=stream,
        )


class ChatCompletion:
    """
    Chat completions API resource.
    """
    
    def __init__(self, client: InfernoClient):
        """
        Initialize the Chat Completions API resource.
        
        Args:
            client: Inferno client
        """
        self.client = client
    
    def create(
        self,
        model: str,
        messages: MessagesInput, # Use the MessagesInput type alias
        temperature: float = 0.7,
        top_p: float = 0.95,
        n: int = 1,
        stream: bool = False,
        stop: Optional[Union[str, List[str]]] = None,
        max_tokens: int = 256,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0,
        logit_bias: Optional[Dict[str, float]] = None,
        user: Optional[str] = None,
        tools: Optional[List[Tool]] = None, # Added tools
        tool_choice: Optional[ToolChoice] = None, # Added tool_choice
    ) -> Union[Dict[str, Any], Iterator[Dict[str, Any]]]:
        """
        Create a chat completion.

        Args:
            model: ID of the model to use.
            messages: A list of messages comprising the conversation so far.
                Can be a list of `Message` dataclasses or a list of dictionaries
                following the OpenAI format.
                The `content` field can be a string or a list of content parts
                (text or image_url) for multimodal input. Image URLs can be
                data URIs (e.g., "data:image/jpeg;base64,...") or standard
                HTTPS URLs.
                Example (using dicts):
                ```python
                messages=[
                    {"role": "user", "content": [
                        {"type": "text", "text": "What's in this image?"},
                        {"type": "image_url", "image_url": {"url": "https://..."}}
                    ]}
                ]
                ```
            temperature: Sampling temperature.
            top_p: Nucleus sampling parameter.
            n: Number of chat completion choices to generate.
            stream: Whether to stream the response
            stop: Up to 4 sequences where the API will stop generating further tokens
            max_tokens: Maximum number of tokens to generate
            presence_penalty: Penalty for token presence.
            frequency_penalty: Penalty for token frequency.
            logit_bias: Modify the likelihood of specified tokens appearing in the completion.
            user: A unique identifier representing your end-user.
            tools: A list of tools the model may call. Currently only functions are supported.
            tool_choice: Controls which tool is called by the model. "none" means no tool call,
                         "auto" lets the model decide, and a specific tool name forces that tool.

        Returns:
            Union[Dict[str, Any], Iterator[Dict[str, Any]]]: Chat completion response
        """
        # Prepare request data
        json_data = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "stream": stream,
            "max_tokens": max_tokens,
        }
        
        # Add optional parameters
        if n != 1:
            json_data["n"] = n
            
        if stop is not None:
            json_data["stop"] = prepare_stop_sequences(stop)
            
        if presence_penalty != 0.0:
            json_data["presence_penalty"] = presence_penalty
            
        if frequency_penalty != 0.0:
            json_data["frequency_penalty"] = frequency_penalty
            
        if logit_bias is not None:
            json_data["logit_bias"] = logit_bias

        if user is not None:
            json_data["user"] = user

        # Add tools and tool_choice if provided
        if tools is not None:
            # Convert Tool dataclasses to dicts if necessary
            json_data["tools"] = [asdict(tool) if is_dataclass(tool) else tool for tool in tools]

        if tool_choice is not None:
             # Convert ToolChoice dataclass to dict if necessary
            json_data["tool_choice"] = asdict(tool_choice) if is_dataclass(tool_choice) else tool_choice

        # Make the request
        return self.client.request(
            method="POST",
            path="/chat/completions",
            json_data=json_data,
            stream=stream,
        )


class Embedding:
    """
    Embeddings API resource.
    """
    
    def __init__(self, client: InfernoClient):
        """
        Initialize the Embeddings API resource.
        
        Args:
            client: Inferno client
        """
        self.client = client
    
    def create(
        self,
        model: str,
        input: Union[str, List[str]],
        user: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Create embeddings.
        
        Args:
            model: ID of the model to use
            input: Input text to embed
            user: A unique identifier representing your end-user
            
        Returns:
            Dict[str, Any]: Embeddings response
        """
        # Prepare request data
        json_data = {
            "model": model,
            "input": input,
        }
        
        # Add optional parameters
        if user is not None:
            json_data["user"] = user
        
        # Make the request
        return cast(Dict[str, Any], self.client.request(
            method="POST",
            path="/embeddings",
            json_data=json_data,
        ))


class Model:
    """
    Models API resource.
    """
    
    def __init__(self, client: InfernoClient):
        """
        Initialize the Models API resource.
        
        Args:
            client: Inferno client
        """
        self.client = client
    
    def list(self) -> Dict[str, Any]:
        """
        List available models.
        
        Returns:
            Dict[str, Any]: List of models
        """
        return cast(Dict[str, Any], self.client.request(
            method="GET",
            path="/models",
        ))
    
    def retrieve(self, model: str) -> Dict[str, Any]:
        """
        Retrieve a model.
        
        Args:
            model: ID of the model to retrieve
            
        Returns:
            Dict[str, Any]: Model details
        """
        return cast(Dict[str, Any], self.client.request(
            method="GET",
            path=f"/models/{model}",
        ))


