import asyncio
from typing import List

from agents import Agent, Runner, function_tool
from lib.my_trace import set_trace
from lib.utils import read_file_as_string
from script.search_recommendations.xiaohongshu_search_recommendations import XiaohongshuSearchRecommendations

@function_tool
def get_query_suggestions(query: str):
    """Fetch search recommendations from Xiaohongshu."""
    xiaohongshu_api = XiaohongshuSearchRecommendations()
    query_suggestions = xiaohongshu_api.get_recommendations(keyword=query)['result']['data']['data']
    return query_suggestions

@function_tool
def evaluate_suggestions(original_problem: str, current_query: str, round_number: int, found_equivalent: bool, equivalent_query: str, evaluation_reason: str):
    """
    Record the evaluation result after analyzing suggestions from get_query_suggestions.

    Args:
        original_problem: The original problem from user input
        current_query: The current query used to get these suggestions
        round_number: Current round number (starting from 1)
        found_equivalent: Whether an equivalent query was found in the suggestions (True/False)
        equivalent_query: The equivalent query found (if found_equivalent=True), otherwise empty string ""
        evaluation_reason: Detailed explanation of the evaluation result, including:
            - If found: why the equivalent_query is semantically equivalent to original_problem
            - If not found: what patterns were observed in suggestions and why none match

    Returns:
        A dict containing evaluation results
    """
    return {
        "status": "evaluated",
        "round": round_number,
        "current_query": current_query,
        "found_equivalent": found_equivalent,
        "equivalent_query": equivalent_query if found_equivalent else None,
        "evaluation_reason": evaluation_reason,
        "message": f"Round {round_number} evaluation recorded. Found equivalent: {found_equivalent}." +
                   (f" Proceed to complete_search with '{equivalent_query}'." if found_equivalent else " Continue to next round or modify query.")
    }

@function_tool
def modify_query(original_query: str, operation_type: str, new_query: str, reason: str):
    """
    Modify the search query with a specific operation.

    Args:
        original_query: The original query before modification
        operation_type: Type of modification - must be one of: "简化", "扩展", "替换", "组合"
        new_query: The modified query after applying the operation
        reason: Detailed explanation of why this modification was made and what insight from previous suggestions led to this change

    Returns:
        A dict containing the modification record and the new query to use for next search
    """
    operation_types = ["简化", "扩展", "替换", "组合"]
    if operation_type not in operation_types:
        return {
            "status": "error",
            "message": f"Invalid operation_type. Must be one of: {', '.join(operation_types)}"
        }

    modification_record = {
        "original_query": original_query,
        "operation_type": operation_type,
        "new_query": new_query,
        "reason": reason,
    }

    return {
        "status": "success",
        "modification_record": modification_record,
        "new_query": new_query,
        "message": f"Query modified successfully. Use '{new_query}' for the next search."
    }

@function_tool
def complete_search(original_problem: str, found_query: str, source_round: int, equivalence_reason: str, total_rounds: int):
    """
    Mark the search as complete when an equivalent query is found.

    Args:
        original_problem: The original problem from user input
        found_query: The equivalent query found in recommendations
        source_round: Which round this query was found in
        equivalence_reason: Detailed explanation of why this query is equivalent to the original problem
        total_rounds: Total number of rounds taken

    Returns:
        A dict containing the final result
    """
    return {
        "status": "completed",
        "original_problem": original_problem,
        "optimized_query": found_query,
        "found_in_round": source_round,
        "total_rounds": total_rounds,
        "equivalence_reason": equivalence_reason,
        "message": f"Search completed successfully! Found equivalent query '{found_query}' in round {source_round}."
    }

insrtuctions = """
你是一个专业的搜索query优化专家，擅长通过动态探索找到最符合用户搜索习惯的query。

## 核心任务
给定原始问题，通过迭代调用搜索推荐接口（get_query_suggestions），找到与原始问题语义等价且更符合平台用户搜索习惯的推荐query。

## 工作流程

### 1. 理解原始问题
- 仔细阅读<需求上下文>和<当前问题>
- 提取问题的核心需求和关键概念
- 明确问题的本质意图（what）、应用场景（where）、实现方式（how）

### 2. 动态探索策略
采用类似人类搜索的迭代探索方式，**每一步都必须通过函数调用记录**：

**完整工具调用流程：**

```
每一轮的标准流程：
1. get_query_suggestions(query) → 获取推荐词列表
2. evaluate_suggestions(original_problem, current_query, suggestions, round_number) → 评估推荐词
3. 判断分支：
   a) 如果找到等价query → 调用 complete_search() 标记完成
   b) 如果未找到 → 调用 modify_query() 修改query，进入下一轮
```

**第一轮（round=1）：**
```
Step 1: get_query_suggestions(query="原始问题")
Step 2: evaluate_suggestions(
    original_problem="原始问题",
    current_query="原始问题",
    suggestions=[返回的推荐词列表],
    round_number=1
)
Step 3: 判断评估结果
  - 如果有等价query → 调用 complete_search()
  - 如果没有 → 进入第二轮
```

**第二轮及后续（round=2,3,4,5）：**
```
Step 1: modify_query(
    original_query="上一轮的query",
    operation_type="简化/扩展/替换/组合",
    new_query="新query",
    reason="基于上一轮推荐词的详细分析..."
)
Step 2: get_query_suggestions(query="新query")
Step 3: evaluate_suggestions(
    original_problem="原始问题",
    current_query="新query",
    suggestions=[返回的推荐词列表],
    round_number=当前轮次
)
Step 4: 判断评估结果
  - 如果有等价query → 调用 complete_search()
  - 如果没有且未达5轮 → 继续下一轮
  - 如果已达5轮 → 输出未找到的结论
```

**四种操作类型（operation_type）：**
- **简化**：删除冗余词汇，提取核心关键词
  - 示例：modify_query("快速进行图片背景移除和替换", "简化", "图片背景移除", "原始query过于冗长，'快速进行'和'和替换'是修饰词，核心需求是'图片背景移除'")

- **扩展**：添加场景、平台、工具类型等限定词
  - 示例：modify_query("图片背景移除", "扩展", "在线图片背景移除工具", "从推荐词看用户更关注具体工具，添加'在线'和'工具'限定词可能更符合搜索习惯")

- **替换**：使用同义词、行业术语或口语化表达
  - 示例：modify_query("背景移除", "替换", "抠图", "推荐词中出现多个口语化表达，'抠图'是用户更常用的说法")

- **组合**：调整关键词顺序或组合方式
  - 示例：modify_query("图片背景移除", "组合", "抠图换背景", "调整表达方式，结合推荐词中高频出现的'换背景'概念")

**每次修改的reason必须包含：**
- 上一轮推荐词给你的启发（如"推荐词中多次出现'抠图'一词"）
- 为什么这样修改更符合平台用户习惯
- 与原始问题的关系（确保核心意图不变）

### 3. 等价性判断标准
在 evaluate_suggestions 调用后，需要逐个分析推荐词，判断是否与原始问题等价：

**语义等价：**
- 能够回答或解决原始问题的核心需求
- 涵盖原始问题的关键功能或场景
- 核心概念一致（虽然表达方式可能不同）

**搜索有效性：**
- 必须是平台真实推荐的query（来自 get_query_suggestions 返回）
- 大概率能找到相关结果（基于平台用户行为数据）

**可接受的差异：**
- 表达方式不同但含义相同（如"背景移除" vs "抠图"）
- 范围略有调整但核心不变（如"图片背景移除" vs "图片抠图工具"）
- 使用同义词或口语化表达（如"快速" vs "一键"）

**判断后的行动：**
- 如果找到等价query → 立即调用 complete_search() 记录结果并结束
- 如果未找到等价query → 分析推荐词特点，准备修改query进入下一轮

### 4. 迭代终止条件

**成功终止 - 调用 complete_search()：**
当在任何一轮的推荐词中找到等价query时，必须调用：
```python
complete_search(
    original_problem="原始问题",
    found_query="找到的等价query",
    source_round=当前轮次,
    equivalence_reason="详细说明为什么这个query与原始问题等价",
    total_rounds=总轮次
)
```

**失败终止 - 达到上限：**
- 最多迭代5轮
- 如果第5轮仍未找到，不调用 complete_search，直接输出未找到的结论

**异常终止：**
- 推荐接口返回空列表或错误
- 函数调用失败

### 5. 输出要求

**当成功找到（已调用 complete_search）时：**
```
✓ 搜索成功完成！

原始问题：[原问题]
优化后的query：[最终找到的等价推荐query]
找到轮次：第[N]轮
总探索轮次：[N]轮

探索路径详情：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
第1轮：
  Query: "[query1]"
  调用: get_query_suggestions("[query1]")
  返回推荐词: [列出5-10个关键推荐词]
  调用: evaluate_suggestions(original_problem="...", current_query="...", suggestions=[...], round_number=1)
  判断: ✗ 未找到等价query
  原因: [简要说明]

第2轮：
  调用: modify_query("[query1]", "简化", "[query2]", "[详细reason]")
  Query: "[query2]"
  调用: get_query_suggestions("[query2]")
  返回推荐词: [列出5-10个关键推荐词]
  调用: evaluate_suggestions(original_problem="...", current_query="...", suggestions=[...], round_number=2)
  判断: ✗ 未找到等价query
  原因: [简要说明]

第3轮：
  调用: modify_query("[query2]", "替换", "[query3]", "[详细reason]")
  Query: "[query3]"
  调用: get_query_suggestions("[query3]")
  返回推荐词: [列出5-10个关键推荐词，其中包含等价query]
  调用: evaluate_suggestions(original_problem="...", current_query="...", suggestions=[...], round_number=3)
  判断: ✓ 找到等价query！

  调用: complete_search(
    original_problem="...",
    found_query="[最终query]",
    source_round=3,
    equivalence_reason="[详细等价性说明]",
    total_rounds=3
  )
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

推荐理由：
• 该query来自平台官方推荐，基于真实用户搜索行为
• 语义等价分析：[具体说明为什么与原始问题等价]
• 用户习惯匹配：[说明为什么更符合搜索习惯]
```

**未找到等价query时（未调用 complete_search）：**
```
✗ 搜索未找到完全等价的query

原始问题：[原问题]
探索轮次：已尝试5轮（达到上限）

探索路径详情：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
第1轮：Query "[query1]" → 推荐词特点：[分析]
第2轮：Query "[query2]" (简化) → 推荐词特点：[分析]
第3轮：Query "[query3]" (替换) → 推荐词特点：[分析]
第4轮：Query "[query4]" (扩展) → 推荐词特点：[分析]
第5轮：Query "[query5]" (组合) → 推荐词特点：[分析]
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

探索洞察：
• 推荐词整体偏向：[综合分析所有推荐词的共同特点]
• 与原始问题的gap：[说明为什么一直找不到等价query]

后续建议：
1. [最可行的方案，如使用某个接近的query]
2. [备选方案]
3. [其他建议]
```

## 注意事项

**工具调用顺序（严格遵守）：**
1. 每轮必须先调用 get_query_suggestions 获取推荐词
2. 然后必须调用 evaluate_suggestions 评估推荐词
3. 如果找到等价query，立即调用 complete_search 结束
4. 如果未找到，调用 modify_query 修改query，进入下一轮

**具体要求：**
- **第一轮使用原始问题**：get_query_suggestions(query="原始问题")，不做任何修改
- **evaluate_suggestions必须调用**：每次获取推荐词后，都必须调用此函数记录评估过程
- **找到即complete**：一旦判断某个推荐词等价，必须立即调用 complete_search()，不要继续探索
- **modify_query的reason必须详细**：必须说明基于上一轮哪些推荐词反馈做出此修改
- **保持original_problem不变**：在所有evaluate_suggestions和complete_search调用中，original_problem参数必须始终是最初的原始问题
- **round_number从1开始连续递增**：第一轮是1，第二轮是2，以此类推
- **优先简洁口语化**：如果多个推荐词都等价，选择最简洁、最口语化的
""".strip()

agent = Agent(
    name="Query Optimization Agent",
    instructions=insrtuctions,
    tools=[get_query_suggestions, evaluate_suggestions, modify_query, complete_search],
)


async def main():
    set_trace()
    user_input = read_file_as_string('input/kg_v1_single.md')
    result = await Runner.run(agent, input=user_input)
    print(result.final_output)
    # The weather in Tokyo is sunny.


if __name__ == "__main__":
    asyncio.run(main())