import json
import uuid
from enum import Enum
from typing import Optional, Literal
from pydantic import BaseModel, model_validator, ConfigDict

from vendi.core.schema import SchemaMixin
from vendi.deployments.schema import DeploymentStatus


class ImageContentUrl(BaseModel):
    url: str


class ImageContent(BaseModel):
    type: Literal["image_url"]
    image_url: ImageContentUrl


class TextContent(BaseModel):
    type: Literal["text"]
    text: str


class LlmMessage(BaseModel):
    role: str
    """The role of the message. Can be `user` or `assistant`."""
    content: str | list[ImageContent | TextContent]
    """The content of the message."""


class CompletionUsage(BaseModel):
    completion_tokens: int
    """Number of tokens in the generated completion."""

    prompt_tokens: int
    """Number of tokens in the prompt."""

    total_tokens: int
    """Total number of tokens used in the request (prompt + completion)."""


class Choice(BaseModel):
    finish_reason: Literal["stop", "length"]
    """The reason the model stopped generating tokens.

    This will be `stop` if the model hit a natural stop point or a provided stop
    sequence, `length` if the maximum number of tokens specified in the request was
    reached.
    """

    index: int
    """The index of the choice in the list of choices."""

    message: LlmMessage
    """A chat completion message generated by the model."""


class ChatCompletion(BaseModel):
    id: str
    """A unique identifier for the chat completion."""

    choices: list[Choice]
    """A list of chat completion choices.

    Can be more than one if `n` is greater than 1.
    """

    created: int
    """The Unix timestamp (in seconds) of when the chat completion was created."""

    model: str
    """The model used for the chat completion."""

    object: Literal["chat.completion"] = "chat.completion"
    """The object type, which is always `chat.completion`."""

    usage: Optional[CompletionUsage] = None
    """Usage statistics for the completion request."""


class VendiCompletionResponse(ChatCompletion):
    id: str | None = None
    """A unique identifier for the completion."""
    elapsed_time: float | None = None
    """The total time taken to generate the completion."""
    provider: str | None = None
    """The name of the provider that served the completion."""
    parameters: any = None
    """The parameters used for the completion."""

    class Config:
        arbitrary_types_allowed = True


class CompletionParams(BaseModel):
    stop: str | list[str] = []
    """The stop sequence(s) to use for the completion."""
    temperature: float = 0.0
    """The sampling temperature to use for the completion."""
    max_tokens: int = 256
    """The maximum number of tokens to generate in the model response."""
    top_k: int = 40
    """The maximum number of tokens to consider for each step of the generation."""
    top_p: float = 0.95
    """The cumulative probability of the top tokens to consider for each step of the generation."""
    frequency_penalty: float = 1.0
    """The frequency penalty to use for the completion."""
    presence_penalty: float = 1.0
    """The presence penalty to use for the completion."""
    n: int = 1
    """The number of completions to generate."""
    json_schema: str | None = None
    """The json schema to enforce on the completion. Used for constrained generation"""
    regex: str | None = None
    """The regex to enforce on the completion. Used for constrained generation"""
    checkpoint: str | None = None
    """The checkpoint to use for the completion. Used for fine-tuned models by vendi"""

    @model_validator(mode="after")
    @classmethod
    def validate_stop(cls, data: "CompletionParams"):
        if isinstance(data.stop, str):
            data.stop = [data.stop]
        elif isinstance(data.stop, list):
            pass
        else:
            data.stop = []
        return data

    @model_validator(mode="after")
    @classmethod
    def validate_schema_and_regex(cls, data: "CompletionParams"):
        if data.json_schema and data.regex:
            raise ValueError(
                "Json schema and regex cannot be used together. Please use on of them"
            )
        if data.json_schema:
            try:
                json.loads(data.json_schema)
            except json.JSONDecodeError:
                raise ValueError(
                    "Passed schema is not valid json. Please check your syntax."
                )
        return data

    model_config = ConfigDict(protected_namespaces=(), extra='allow')


class ModelParameters(CompletionParams):
    model: str
    """The model to use for the completion - in the format <provider>/<model-name>"""


class CompletionRequest(ModelParameters):
    request_id: str | None = None
    """A unique identifier for the completion request. Can be used to track the request."""
    messages: list[LlmMessage]
    """The prompt messages to use for the completion."""


class BatchInferenceStatus(str, Enum):
    CREATING = "creating"
    """The batch inference is being created."""
    PENDING = "pending"
    """The batch job is created and waits for execution"""
    RUNNING = "running"
    """The batch job is running."""
    FAILED = "failed"
    """The batch job failed."""
    COMPLETED = "completed"
    """The batch job successfully completed."""


class BatchInference(SchemaMixin):
    model_parameters: list[ModelParameters]
    """The model parameters for the batch inference."""
    dataset_id: uuid.UUID
    """The dataset id for the batch inference."""
    output_dataset_id: uuid.UUID
    """The outputs of the batch job are saved in this dataset id. This should be used to retrieve the results."""
    status: BatchInferenceStatus
    """The status of the batch inference."""
    total_chunks: int | None = None
    """The batch job is split into chunks. This is the total number of chunks."""
    total_records: int | None = None
    """The total number of records in the batch inference."""
    completed_chunks: int | None = None
    """The number of completed chunks. This is updated as the batch job progresses"""
    job_name: str | None = None
    """The name of the batch job."""

    model_config = ConfigDict(
        protected_namespaces=(),
    )


class Endpoint(SchemaMixin):
    model_id: str
    """The model identifier served in this endpoint"""
    provider: str
    """The provider of the model (vendi/openai/google..)"""
    adapter_config: str | None = None
    """The lora adapter configuration for the model"""
    short_description: str | None = None
    """A description of the model"""
    full_description: str | None = None
    """A full description of the model"""
    docs: str | None = None
    """A doc link about the model"""
    status: DeploymentStatus | None = None
    """The status of the endpoint. If the endpoint is not ready, the status will be `not_ready`"""
    is_configured: bool | None = None
    """If the endpoint is configured by the user or not"""
