from typing import Any, Dict, List, Literal, Optional

from openai.types.chat.chat_completion import ChatCompletion
from pydantic import BaseModel, Field

type OptionalParams = Dict[str, Any]


class Response(BaseModel):
    status_code: int
    request_id: str = Field(description="ID of the request, generated by the server")
    body: Optional[ChatCompletion] = None


class BatchResponse(BaseModel):
    id: str = Field(description="The ID of the batch request")
    custom_id: str
    response: Optional[Response] = None
    error: Optional[Any] = None


class Body(BaseModel):
    messages: List[Dict[str, Any]]
    model: str
    temperature: float
    max_tokens: int

    class Config:
        extra = "allow"


class OpenAIBatch(BaseModel):
    custom_id: str
    method: Literal["POST"] = "POST"
    url: str = "/v1/chat/completions"
    body: Body


class AnthropicBatch(BaseModel):
    custom_id: str
    params: Body


class OutputModel(BaseModel):
    custom_id: str
    type: str
    model: str
    response: str
    input_tokens: int
    output_tokens: int


class BatchConfig(BaseModel):
    format: Literal["openai", "anthropic"]

    model: str
    temperature: float
    max_tokens: int
    n_answers: int = 1
    system_message: Optional[str] = None
    response_model: Optional[dict] = None
    kwargs: Optional[OptionalParams] = None


class Parameters(BaseModel):
    model: str
    temperature: float
    max_tokens: int

    class Config:
        extra = "allow"


class Config(BaseModel):
    format: Literal["openai", "anthropic"]
    params: Parameters
    n_answers: int = 1
    system_message: Optional[str] = None
    json_schema: Optional[dict] = None


class Question(BaseModel):
    question_id: str
    question: str
    image_path: Optional[str] = None
