from typing import List, Union, Optional, Literal, Dict, Any
from pydantic import BaseModel, Field
from just_agents.data_classes import Message, ToolCall, Role
from just_agents.llm_options import ModelOptions

### OpenAI API compatible models ###

class ChatCompletionRequest(ModelOptions):
    messages: List[Message] = Field(..., examples=[[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What drug interactions of rapamycin are you aware of? What are these interactions ?"}
    ]])
    n: Optional[int] = Field(1, ge=1)
    stream: Optional[bool] = Field(default=False, examples=[True])
    stop: Optional[Union[str, List[str]]] = Field(default=None, examples=[None, "[STOP]"])
    max_tokens: Optional[int] = Field(None, ge=1)
    logit_bias: Optional[dict] = Field(None, examples=[None])
    user: Optional[str] = Field(None, examples=[None])


class ResponseMessage(Message):
    refusal: Optional[str] = Field(None, description="The refusal message generated by the model.")
    audio: Optional[dict] = Field(None, description="Audio response from the model if the audio output modality is requested.")
    function_call: Optional[dict] =  Field(None, description="Deprecated and replaced by `tool_calls`")
    tool_calls: Optional[List[ToolCall]] = Field(None, description="The tool calls generated by the model, such as function calls.")


class Logprob(BaseModel):
    token: str = Field(..., description="The token.")
    bytes: Optional[List[int]] = Field(None,description="""
    A list of integers representing the UTF-8 bytes representation of the token.
    Useful in instances where characters are represented by multiple tokens and their byte representations must be 
    combined to generate the correct text representation. Can be `null` if there is no bytes representation for the token."""
    )
    logprob: float = Field(-9999.0,description="""
    The log probability of this token, if it is within the top 20 most likely tokens.
    Otherwise, the value `-9999.0` is used to signify that the token is very unlikely."""
    )

class TokenLogprob(Logprob):
    top_logprobs: List[Logprob] = Field(default_factory=list,description="List of the most likely tokens and their log probability, at this token position.")

class ChoiceLogprobs(BaseModel):
    content: Optional[List[TokenLogprob]] = Field(None,description="A list of message content tokens with log probability information.")
    refusal: Optional[List[TokenLogprob]] = Field(None,description="A list of message refusal tokens with log probability information.")

class ChatCompletionChoice(BaseModel):
    index: int = Field(0,ge=0, description="The index of the choice in the list of choices.")
    logprobs: Optional[ChoiceLogprobs] = Field( None,description="Log probability information for the choice.")
    finish_reason: Optional[Literal["stop", "length", "tool_calls", "content_filter", "function_call"]] = Field(
        ...,
        description="The reason the model stopped generating tokens."
    )
    message : Optional[ResponseMessage]

class ChatCompletionChoiceChunk(ChatCompletionChoice):
    delta: ResponseMessage = Field(default=None)
    message: Optional[ResponseMessage] = Field(default=None, exclude=True) # removed, delta is used instead

class CompletionTokensDetails(BaseModel):
    accepted_prediction_tokens: Optional[int] = Field(None,
        description="When using Predicted Outputs, the number of tokens in the prediction that appeared in the completion.")
    audio_tokens: Optional[int] = Field(None, description="""Audio input tokens generated by the model.""")
    reasoning_tokens: Optional[int] = Field(None, description="""Tokens generated by the model for reasoning.""")

    rejected_prediction_tokens: Optional[int] = Field(None, description="""
    When using Predicted Outputs, the number of tokens in the prediction that did not appear in the completion. 
    However, like reasoning tokens, these tokens are still counted in the total completion tokens for purposes of billing, output, and context window limits.
    """)

class PromptTokensDetails(BaseModel):
    audio_tokens: Optional[int] = Field(None, description="""Audio input tokens present in the prompt.""")
    cached_tokens: Optional[int] = Field(None, description="""Cached tokens present in the prompt.""")

class ChatCompletionUsage(BaseModel):
    prompt_tokens: int = Field(default=0, description="""Number of tokens in the prompt.""")
    completion_tokens: int = Field(default=0, description="""Number of tokens in the generated completion.""")
    total_tokens: int = Field(default=0, description="""Total number of tokens used in the request (prompt + completion).""")
    completion_tokens_details: Optional[CompletionTokensDetails] = Field(None, description="""Breakdown of tokens used in a completion.""")
    prompt_tokens_details: Optional[PromptTokensDetails] = Field(None, description="""Breakdown of tokens used in the prompt.""")

class ChatCompletionResponse(BaseModel):
    id: str = Field(..., description="A unique identifier for the chat completion.")
    object: Literal["chat.completion", "chat.completion.chunk"] = Field(...,description="The object type")
    created: Union[int,float] = Field(..., description="The Unix timestamp (in seconds) of when the chat completion was created.")
    model: str = Field(..., description="The model used for the chat completion.")
    choices: List[ChatCompletionChoice] = Field(...,description="A list of chat completion choices. Can be more than one if `n` is greater than 1.")
    system_fingerprint: Optional[str] = Field(None, description="This fingerprint represents the backend configuration that the model runs with.")
    service_tier: Optional[str] = Field(None, exclude=True, description="This field is only included if the `service_tier` parameter is specified in the request.")
    usage: Optional[ChatCompletionUsage] = Field(default=None, description= "Usage statistics for the completion request.")

class ChatCompletionChunkResponse(ChatCompletionResponse):
    choices: List[ChatCompletionChoiceChunk] = Field(...,description="A list of chunked choices. Can be more than one if `n` is greater than 1.")

# Model definitions for API responses

class Model(BaseModel):
    """Represents a model in the API response, following OpenAI's model format"""
    id: str = Field(
        description="Unique identifier for the model",
        examples=["gpt-3.5-turbo", "claude-2", "llama-2-7b"]
    )
    created: int = Field(
        description="Unix timestamp of model creation",
        examples=[1677649963]
    )
    object: str = Field(
        default="model", 
        description="Object type in the API response",
        examples=["model"]
    )
    owned_by: str = Field(
        default="organization", 
        description="Entity that owns the model",
        examples=["organization", "anthropic", "openai"]
    )
    permission: Optional[List[Dict[str, Any]]] = Field(
        default=None, 
        description="Array of permissions associated with the model",
        examples=[[{"id": "modelperm-123", "type": "access"}]]
    )
    root: Optional[str] = Field(
        default=None, 
        description="The base model this model is derived from",
        examples=["gpt-3.5"]
    )
    parent: Optional[str] = Field(
        default=None, 
        description="The immediate parent model in a model hierarchy",
        examples=["gpt-3.5-turbo"]
    )

class ModelList(BaseModel):
    """Container for list of models in API response, following OpenAI's model list format."""
    object: str = Field(
        default="list", 
        description="Type of object in the API response",
        examples=["list"]
    )
    data: List[Model] = Field(
        ..., 
        description="List of available models",
        min_length=0,
        examples=[
            [
                {
                    "id": "gpt-3.5-turbo", 
                    "created": 1677649963, 
                    "object": "model", 
                    "owned_by": "organization"
                }
            ]
        ]
    )


class ErrorResponse(BaseModel):
    class ErrorDetails(BaseModel):
        message: str = Field(...)
        type: str = Field("server_error")
        code: str = Field("internal_server_error")
    error: ErrorDetails
