from typing import Any, Optional

from pydantic import BaseModel, Field


class Image(BaseModel):
    content_type: str
    b64encoded: str


class Answer(BaseModel):
    answer: str
    original_question_uuid: Optional[str]
    actual_question_uuid: Optional[str]
    module: str
    agent_path: str


class Step(BaseModel):
    original_question_uuid: Optional[str]
    actual_question_uuid: Optional[str]
    module: str
    title: str
    value: Optional[str] = None
    agent_path: str
    reason: Optional[str] = None
    timeit: float
    input_nuclia_tokens: Optional[float]
    output_nuclia_tokens: Optional[float]
    error: Optional[str] = None

    def __str__(self) -> str:
        return (
            f"({self.timeit}) {self.module}: {self.title} - {self.value} : "
            f"{self.reason} ({self.input_nuclia_tokens}:{self.output_nuclia_tokens})"
        )


class Chunk(BaseModel):
    chunk_id: str
    title: Optional[str] = None
    source: Optional[str] = None
    text: str
    labels: list[str] = Field(default_factory=list)
    url: list[str] = Field(default_factory=list)
    metadata: Optional[dict[str, Any]] = None


class Prompt(BaseModel):
    prompt: str
    resources: list[str] = Field(default_factory=list)
    links: list[str] = Field(default_factory=list)
    description: Optional[str] = None


class Context(BaseModel):
    original_question_uuid: Optional[str]
    actual_question_uuid: Optional[str]
    question: str
    chunks: list[Chunk] = Field(default_factory=list)
    images: dict[str, Image] = Field(default_factory=dict)
    prompts: list[Prompt] = Field(default_factory=list)
    structured: list[str] = Field(default_factory=list)
    source: str
    agent: str
    # XXX: This is not actually a summary, but an answer attempt for now!
    summary: str = Field(
        default="",
        description="Partial or full answer to the question,"
        " generated by the context validation step inside a context agent.",
    )
    agent_id: str = ""
    title: Optional[str] = None
    missing: Optional[str] = None
    citations: Optional[list[str]] = Field(
        default=None,
        description="list of chunk IDs that were considered relevant in the context validation step.",
    )
