from pydantic import BaseModel
from typing import List, Dict, Union


class QuestionAnsweringSystemPrompt(BaseModel):
    system_prompt: str = (
        """Answer the user's question based on the given information. Do NOT make up any facts."""
    )
    context_documents: List[str] = []
    user_query: str = ""

    @property
    def compiled_system_prompt(self):
        return self.system_prompt + "\n\n".join(self.context_documents)

    @property
    def messages(self):
        return [
            {"role": "system", "content": self.compiled_system_prompt},
            {"role": "user", "content": self.user_query},
        ]

    def compile(self):
        return self.compiled_system_prompt


class SummarisationSystemPrompt(BaseModel):
    text_to_summarise: str
    n_sentences: int

    @property
    def system_prompt(self) -> str:
        return f"""Summarise the user's text into the smallest number of points possible. It must NOT be more than {str(self.n_sentences)} sentences.
If the text contains comments posted by people, ignore such comments.
Ensure that your summary contains the technical points and language used in the original text.
"""

    @property
    def messages(self):
        return [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": self.text_to_summarise},
        ]

    def compile(self):
        return self.system_prompt


class TableDescriptionGenerationSystemPrompt(BaseModel):
    table_to_describe: str
    n_sentences: int

    @property
    def system_prompt(self) -> str:
        return f"""Describe this table in the smallest number of points possible. It must NOT be more than {str(self.n_sentences)} sentences.
In your description, mention the various columns and the type of data populated across rows.
Ensure that your summary contains the technical points and language used in the original text.
"""

    @property
    def messages(self):
        return [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": self.table_to_describe},
        ]

    def compile(self):
        return self.system_prompt


class QueryRoutingSystemPrompt(BaseModel):
    system_prompt: str = (
        """Return the most suitable option that this request is related to.
Only return the name of the option, and nothing else, or your response will not be parsed correctly!"""
    )
    routes: List[str] = []
    query: str = ""

    @property
    def compiled_system_prompt(self):
        return f"{self.system_prompt}\n- " + "\n- ".join(self.routes)

    @property
    def messages(self):
        return [
            {"role": "system", "content": self.compiled_system_prompt},
            {"role": "user", "content": self.query},
        ]

    def compile(self):
        return self.compiled_system_prompt
