from typing import List

from app.api.routers.models import Message
from llama_index.core.prompts import PromptTemplate
from llama_index.core.settings import Settings
from pydantic import BaseModel

NEXT_QUESTIONS_SUGGESTION_PROMPT = PromptTemplate(
    "你是一个乐于助人的助手！你的任务是对用户可能会问的下一个问题给出建议。 "
    "\n这是对话历史记录"
    "\n---------------------\n{conversation}\n---------------------"
    "考虑到对话历史记录，仅限于现在知识库已有内容, 请给我 $number_of_questions 个你接下来可能会问题的问题！"
)
N_QUESTION_TO_GENERATE = 3


class NextQuestions(BaseModel):
    """A list of questions that user might ask next"""

    questions: List[str]


class NextQuestionSuggestion:
    @staticmethod
    async def suggest_next_questions(
        messages: List[Message],
        number_of_questions: int = N_QUESTION_TO_GENERATE,
    ) -> List[str]:
        # Reduce the cost by only using the last two messages
        last_user_message = None
        last_assistant_message = None
        for message in reversed(messages):
            if message.role == "user":
                last_user_message = f"User: {message.content}"
            elif message.role == "assistant":
                last_assistant_message = f"Assistant: {message.content}"
            if last_user_message and last_assistant_message:
                break
        conversation: str = f"{last_user_message}\n{last_assistant_message}"

        output: NextQuestions = await Settings.llm.astructured_predict(
            NextQuestions,
            prompt=NEXT_QUESTIONS_SUGGESTION_PROMPT,
            conversation=conversation,
            nun_questions=number_of_questions,
        )

        return output.questions