from __future__ import annotations

from typing import Any

from pydantic import BaseModel, field_validator, model_validator
from typing_extensions import Self

from ragelo.logger import logger
from ragelo.types.results import EvaluatorResult


class ChatMessage(BaseModel):
    sender: str
    content: str

    def __str__(self) -> str:
        return f"{self.sender}: {self.content}"


class Evaluable(BaseModel):
    """A base class for objects that can be evaluated. Either a retrieved document or an agent answer.
    Args:
        evaluation Optional[EvaluatorResult]: The result of the evaluation.
        metadata Optional[dict[str, Any]]: Metadata that can be templated in the prompt.
    """

    qid: str
    evaluation: EvaluatorResult | None = None
    metadata: dict[str, Any] | None = None

    @field_validator("qid", mode="before")
    def qid_into_string(cls, v):
        if not isinstance(v, str):
            try:
                v = str(v)
            except ValueError:
                raise ValueError("qid must be a string or convertible to a string")
        return v

    def add_metadata(self, metadata: dict[str, Any] | None):
        if not metadata:
            return
        if self.metadata is None:
            self.metadata = {}
        for k in metadata:
            if k in self.metadata:
                logger.warning(
                    f"Metadata {k} for {self.__class__.__name__}"
                    " is being overwritten!\n"
                    f"Old metadata: {self.metadata[k]}\n"
                    f"New metadata: {metadata[k]}\n"
                )
            self.metadata[k] = metadata[k]


class Document(Evaluable):
    """A document retrieved by an agent in response to a query.
    Args:
        did str: The document ID.
        text str: The text of the document.
        retrieved_by dict[str, float]: If a document was retrieved by multiple agents,
            the score attributed by each agent for this document.
    """

    did: str
    text: str
    retrieved_by: dict[str, float] = {}

    @field_validator("did", mode="before")
    def did_into_string(cls, v):
        if not isinstance(v, str):
            try:
                v = str(v)
            except ValueError:
                raise ValueError("did must be a string or convertible to a string")
        return v

    def add_retrieved_by(self, agent: str, score: float | None = None, force: bool = False, exist_ok: bool = False):
        """Adds the score of an agent that retrieved the document."""
        if agent in self.retrieved_by and not force:
            if not exist_ok:
                logger.info(f"Document with did {self.did} already retrieved by agent {agent}")
            return
        if score is None:
            score = 1.0
        self.retrieved_by[agent] = score

    def __str__(self) -> str:
        if len(self.text) > 100:
            return f"{self.did}: {self.text[:100]}..."
        return f"{self.did}: {self.text}"

    @classmethod
    def assemble_document(
        cls,
        document: Self | str,
        qid: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Self:
        """Assembles a Document object from a string or a Document object."""
        if isinstance(document, str):
            if qid is None:
                raise ValueError("qid must be provided if document is a string")
            did = "<no_did>"
            if metadata:
                valid_id_fields = ["did", "doc_id", "document_id", "id", "_id"]
                valid_id_fields = [f for f in valid_id_fields if f in metadata]
                if valid_id_fields:
                    did = metadata[valid_id_fields[0]]
            document = cls(qid=qid, did=did, text=document)
        document.add_metadata(metadata)
        return document

    @staticmethod
    def assemble_documents(
        documents: list,
        qid: str,
        metadata: list[dict[str, Any]] | list[None] | None = None,
    ) -> dict[str, "Document"]:
        """Assembles a list of Document objects from a list of strings or Document objects."""
        assembled_docs: dict[str, Document] = {}
        if metadata and len(documents) != len(metadata):
            raise ValueError("The number of documents and document metadata do not match")
        if not metadata:
            metadata = [None] * len(documents)

        for idx, (doc, m) in enumerate(zip(documents, metadata)):
            doc_obj = Document.assemble_document(doc, qid, m)
            if doc_obj.did == "<no_did>":
                doc_obj.did = f"doc_{idx}"
            assembled_docs[doc_obj.did] = doc_obj
        return assembled_docs


class AgentAnswer(Evaluable):
    """An answer generated by an agent in response to a query.
    Args:
        agent str: The agent that provided the answer.
        text str: The text of the answer.
        conversation Optional[list[ChatMessage]]: The conversation between the user and the agent.
    """

    agent: str
    text: str | None = None
    conversation: list[ChatMessage] | None = None

    @model_validator(mode="before")
    @classmethod
    def one_of_text_or_conversation(cls, values):
        text = values.get("text")
        conversation = values.get("conversation")
        if text is None and conversation is None:
            raise ValueError("Either text or conversation must be provided")

        if text is not None and conversation is not None:
            raise ValueError("Only one of text or conversation must be provided")
        return values

    @classmethod
    def assemble_answer(
        cls,
        answer: Self | str,
        qid: str,
        agent: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> "AgentAnswer":
        """Assembles an AgentAnswer object from a string or an AgentAnswer object."""
        if isinstance(answer, str):
            if agent is None:
                raise ValueError("agent must be provided if answer is a string")
            answer = cls(agent=agent, qid=qid, text=answer)
        answer.add_metadata(metadata)
        return answer


class PairwiseGame(Evaluable):
    """A game to be played between two agent answers.
    Args:
        agent_a_answer AgentAnswer: The answer provided by the first agent.
        agent_b_answer AgentAnswer: The answer provided by the second agent.
    """

    agent_a_answer: AgentAnswer
    agent_b_answer: AgentAnswer
