from typing import List, Optional, Union

from pydantic import BaseModel, ConfigDict, Field, field_validator

from lion.integrations.perplexity_.api_endpoints.data_models import (
    Citation,
    PerplexityEndpointResponseBody,
    Usage,
)


class Message(BaseModel):
    """A message in the chat completion response."""

    role: str = Field(description="The role of the message author.")
    content: str = Field(description="The content of the message.")

    model_config = ConfigDict(extra="forbid")

    @field_validator("role")
    @classmethod
    def validate_role(cls, role: str) -> str:
        valid_roles = ["assistant"]  # Response messages are always from assistant
        if role not in valid_roles:
            raise ValueError(f"Role must be one of {valid_roles}")
        return role

    @field_validator("content")
    @classmethod
    def validate_content(cls, content: str) -> str:
        if not content or not content.strip():
            raise ValueError("Content cannot be empty")
        return content


class Choice(BaseModel):
    """A completion choice returned by the model."""

    index: int = Field(description="The index of this choice in the list of choices.")
    message: Message = Field(description="The message generated by the model.")
    finish_reason: str | None = Field(
        description="The reason the model stopped generating tokens. Either 'stop' or 'length'."
    )

    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)

    @field_validator("finish_reason")
    @classmethod
    def validate_finish_reason(cls, finish_reason: str | None) -> str | None:
        if finish_reason is not None:
            valid_reasons = ["stop", "length"]
            if finish_reason not in valid_reasons:
                raise ValueError(f"Finish reason must be one of {valid_reasons}")
        return finish_reason


class RelatedQuestion(BaseModel):
    """A related question suggested by the model."""

    text: str = Field(description="The text of the related question.")

    model_config = ConfigDict(extra="forbid")

    @field_validator("text")
    @classmethod
    def validate_text(cls, text: str) -> str:
        if not text or not text.strip():
            raise ValueError("Text cannot be empty")
        return text


class PerplexityChatCompletionResponseBody(PerplexityEndpointResponseBody):
    """Response body for chat completion requests."""

    id: str = Field(description="A unique identifier for this completion.")
    model: str = Field(description="The model used for completion.")
    object: str = Field(description="The object type, always 'chat.completion'.")
    created: int = Field(
        description="Unix timestamp of when the completion was created."
    )
    choices: list[Choice] = Field(description="The list of completion choices.")
    usage: Usage = Field(description="Token usage information for this request.")
    citations: list[Citation | dict | str] | None = Field(
        default=None,
        description="Citations for the generated answer. Can be Citation objects, dictionaries, or URL strings.",
    )
    related_questions: list[RelatedQuestion] | None = Field(
        default=None, description="Related questions suggested by the model."
    )

    model_config = ConfigDict(extra="forbid")

    @field_validator("object")
    @classmethod
    def validate_object(cls, object_type: str) -> str:
        if object_type != "chat.completion":
            raise ValueError("Object type must be 'chat.completion'")
        return object_type

    @field_validator("choices")
    @classmethod
    def validate_choices(cls, choices: list[Choice]) -> list[Choice]:
        if not choices:
            raise ValueError("At least one choice is required")
        return choices

    @field_validator("citations")
    @classmethod
    def validate_citations(
        cls, citations: list[Citation | dict | str] | None
    ) -> list[Citation | dict | str] | None:
        if citations is not None:
            if len(citations) > 10:  # Arbitrary limit for safety
                raise ValueError("Too many citations")
            # Convert string citations to Citation objects
            processed_citations = []
            for citation in citations:
                if isinstance(citation, str):
                    processed_citations.append({"url": citation})
                else:
                    processed_citations.append(citation)
            return processed_citations
        return citations

    @field_validator("related_questions")
    @classmethod
    def validate_related_questions(
        cls, questions: list[RelatedQuestion] | None
    ) -> list[RelatedQuestion] | None:
        if questions is not None and len(questions) > 10:  # Arbitrary limit for safety
            raise ValueError("Too many related questions")
        return questions
