"""
SmartWebSearch.QueryStorm
~~~~~~~~~~~~

This module implements the query brainstorm for the web searching module.
"""

# Import the required modules
from typing import Any
import requests
import datetime
from SmartWebSearch.AIModel import AIModel
from langdetect import detect

# QueryStorm Class
class QueryStorm:
    """
    A class for query brainstorming.
    """

    def __init__(self, ai_model: AIModel) -> None:
        """
        Initialize the QueryStorm object.

        Args:
            ai_model (AIModel): The AIModel object.

        Returns:
            None
        """

        # Set the attributes of the QueryStorm object
        self.ai_model: AIModel = ai_model

    def decompose_tasks_with_prompt(self, u_prompt: str) -> list[str]:
        """
        Decompose task prompts based on the user prompt.

        Args:
            u_prompt (str): The user prompt.

        Returns:
            list[str]: The generated task prompts.
        """

        prompt: str = """你是一个专业的搜索任务分解助手。你的工作是将用户给出的包含多个搜索需求的提示词，分解成多个独立的搜索任务提示词。

        任务描述
        用户会输入一段提示词{prompt}，其中可能包含一个或多个搜索请求，每个请求针对不同的主体（如人物、团体、事物等）。你需要识别出每个独立的主体，并为每个主体生成一个对应的搜索任务提示词。如果一个主体下用户要求搜索多个方面的内容（如简介、成员、专辑等），则将这些内容整合到同一个任务提示词中。

        输出格式
        - 每个任务提示词的格式为：“搜索关于[主体]的[具体搜索范围]”
        - 如果有多个任务，请用“&&”将每个任务提示词连接起来。
        - 如果只有一个任务，直接输出该任务提示词，不加“&&”。
        - 输出时不要包含任何其他解释或文字，只输出格式化后的任务提示词字符串。
        - 输出的语言必须为“{lang}”，以确保输出语言与用户提示词的语言一致。

        示例
        示例1：
        用户输入：搜索关于IVE这个KPOP组合的资料，关于IVE的简介、成员、专辑和歌曲。此外，帮我搜索关于ATEEZ这个KPOP组合的资料，关于ATEEZ的成员、专辑和歌曲。
        你的输出：搜索IVE这个KPOP组合的简介、成员、专辑和歌曲&&搜索ATEEZ这个KPOP组合的成员、专辑和歌曲

        示例2：
        用户输入：搜索关于IVE这个KPOP组合的资料，关于IVE的简介、成员、专辑和歌曲。此外，帮我搜索关于ATEEZ这个KPOP组合的资料，关于ATEEZ的成员、专辑和歌曲。另外，还有关于HoneyWorks这个日本组合的资料
        你的输出：搜索IVE这个KPOP组合的简介、成员、专辑和歌曲&&搜索ATEEZ这个KPOP组合的成员、专辑和歌曲&&搜索HoneyWorks这个日本组合

        示例3：
        用户输入：搜索关于IVE这个KPOP组合的资料
        你的输出：搜索IVE这个KPOP组合

        示例4：
        用户输入：Search about latest technology news and latest events in the world
        你的输出：Search about latest technology news&&Search about latest events in the world

        示例5：
        用户输入：What is RAG? What advantages does RAG have? What can RAG do?
        你的输出：Search about RAG, its advantages and what can it do
        *解释：虽然用户提问了三个问题，但是都是围绕着同一主题RAG，所以应该将这个三个问题合并成一个任务提示词。

        注意事项
        - 仔细阅读用户输入，识别出所有不同的搜索主体。
        - 每个主体对应一个任务，即使对同一个主体有多个描述，也只生成一个任务提示词。
        - 任务提示词应准确反映用户要求搜索的具体内容，尽量保留用户的原词。
        - 如果用户对某个主体没有指定具体范围，则使用“资料”概括。
        - 确保输出格式严格遵循要求：多个任务用“&&”分隔。
        
        请严格按照上述格式和示例执行。"""

        # Decompose the prompt into task prompts
        res: dict[str, Any] = self.ai_model.send_request(
            [
                {
                    "role": "user",
                    "content": prompt.format(prompt = u_prompt, lang = detect(u_prompt))
                }
            ]
        )

        # Return the decomposed task prompts
        return [ i.strip() for i in res["choices"][0]["message"]["content"].split("&&") ]

    def storm_with_summary(self, u_prompt: str, summary: str) -> list[str]:
        """
        Generate auxiliary queries based on the prompt and summary of the search results.

        Args:
            u_prompt (str): The prompt.
            summary (str): The summary of the search results.

        Returns:
            list[str]: The generated queries. (Auxiliary Queries)
        """

        prompt: str = """你是一个智能搜索助手，专门负责分析用户的搜索意图并生成扩展的搜索辅助关键词。

        任务描述：
        根据用户提供的提示詞“{prompt}”和提示詞相關关键词的搜索结果总结“{summary}”，首先判断用户想搜索的内容的核心类型（例如，概念定义、工具使用、历史背景、技术原理等），然后基于这个类型，延展出更多相关的搜索辅助关键词。这些关键词应帮助用户进一步精确搜索，获取更具体、更深入的信息。

        今天的日期时间是“{datetime}”，如果你需要用到这个日期时间用来搜索最近的内容，请在搜索辅助关键词中包含它。
        
        输出格式要求：
        - 输出3至12个搜索辅助关键词即可，不必太多，也不能太少。
        - 仅输出搜索辅助关键词，不包含任何其他文字或解释。
        - 每个搜索辅助关键词之间用一个空格“ ”隔开。
        - 如果搜索辅助关键词内包含多个单词，请用加号“+”连接，不要使用空格或其他分隔符。
        - 输出的语言必须为“{lang}”，以确保输出语言与用户提示词的语言一致。

        示例：
        输入：
        prompt = 什麽是三角函数？
        summary = 三角函数是数学很常见的一类关于角度的函数。三角函数将直角三角形的内角和它的两边的比值相关联，亦可以用单位圆的各种有关线段的长短的等价来定义。三角函数在研究三角形和圆形等几何形状的性质时有着重要的作用，亦是研究振动、波、天体运动和各种周期性现象的基础数学工具。在数学分析上，三角函数亦定义为无穷级数或特定微分方程式的解，允许它们的取值扩展到任意实数值，甚至是复数值。
        输出：
        definitions purposes general+formulas

        请严格按照上述格式和示例执行。"""

        # Generate queries based on the summary of the search results
        res: dict[str, Any] = self.ai_model.send_request(
            [
                {
                    "role": "user",
                    "content": prompt.format(prompt = u_prompt, summary = summary, datetime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), lang = detect(u_prompt))
                }
            ]
        )

        # Return the generated queries
        return res["choices"][0]["message"]["content"].split(" ")
    
    def storm_with_prompt(self, u_prompt: str) -> list[str]:
        """
        Generate a query based on the prompt.

        Args:
            u_prompt (str): The user prompt.

        Returns:
            list[str]: The generated queries. (Main Queries, Auxiliary Queries)
        """

        prompt: str = """你是一个专业的搜索关键词优化助手，擅长分析用户的查询意图，并提取精准的网络搜索关键词。

        任务描述：
        根据用户提供的提示词 {prompt}，判断用户想要搜索的核心内容，并列出网络搜索关键词。

        今天的日期时间是“{datetime}”，如果你需要用到这个日期时间用来搜索最近的内容，请在搜索辅助关键词中包含它。

        关键词生成规则：
        - 主要关键词：准确反映搜索主题的核心概念，必须保留。
        - 辅助关键词：围绕主要关键词，涵盖用户明确提及需要搜索的具体内容（例如定义、类型、用途、原理、历史、相关公式或操作方法等）。
        - 输出长度规则：
        - 如果用户仅提出核心搜索主题，未明确要求搜索任何具体内容，则只输出主要关键词。
        - 如果提示词中提到了时间点，如年份、月份和日期，应该把这些时间点以及相关的关键词包含在主要关键词中，而不是把时间点单独作为主要关键词。
        - 如果用户在提示词中明确指明需要搜索与核心主题相关的具体内容，则输出「主要关键词 + 辅助关键词」，且辅助关键词数量不超过5个。
        - 不要生成用戶沒有指明需要搜索与核心主题相关的具体内容的搜索辅助关键词。
        - 关键词语言：关键词语言必须为“{lang}”，以确保输出语言与用户提示词的语言一致。
        
        输出格式要求：
        - 仅输出关键词本身，不包含任何解释或附加文字。
        - 关键词之间用一个空格「 」分隔。
        - 如果一个关键词由多个词语组成，请用加号「+」连接（例如：artificial+intelligence 或 人工智能+应用），以确保该关键词在搜索时被视为一个整体。
        - 第一个关键词始终为主要关键词，若存在辅助关键词则依次排列其后（最多5个）。
        
        示例：
        - 用户提示词：What is trigonometry?（未指定具体内容）
        输出：trigonometric+functions
        - 用户提示词：I want to learn about the definitions and formulas of trigonometry（明确指定内容）
        输出：trigonometric+functions definitions formulas applications （此处“definitions”和“formulas”是用户明确提到的，因此优先纳入）
        - 用户提示词：搜索2026年各KPOP組合的專輯和歌曲（明确指定时间点）
        输出：2026+kpop 專輯 歌曲
        - 用户提示词：搜索2025年的科技新聞（明确指定时间点）
        输出：2025+科技+新聞
        
        请严格遵循上述格式，确保输出的关键词准确、简洁，能够帮助用户进行高效的网络搜索。"""

        # Generate a query based on the prompt
        res: dict[str, Any] = self.ai_model.send_request(
            [
                {
                    "role": "user",
                    "content": prompt.format(prompt = u_prompt, datetime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), lang = detect(u_prompt)),
                }
            ]
        )

        # Return the generated queries
        return res["choices"][0]["message"]["content"].split(" ")