import asyncio
import json
import os
import argparse
from datetime import datetime

from agents import Agent, Runner, function_tool, AgentOutputSchema
from lib.my_trace import set_trace
from typing import Literal
from pydantic import BaseModel, Field


from lib.utils import read_file_as_string
from script.search_recommendations.xiaohongshu_search_recommendations import XiaohongshuSearchRecommendations

from agents import Agent, RunContextWrapper, Runner, function_tool

from pydantic import BaseModel, Field


class RunContext(BaseModel):
    version: str = Field(..., description="当前运行的脚本版本（文件名）")
    input_files: dict[str, str] = Field(..., description="输入文件路径映射，如 {'context_file': '...', 'q_file': '...'}")
    q_with_context: str
    q_context: str
    q: str
    log_url: str
    log_dir: str
    # 问题标注结果 - 直接用字符串记录
    question_annotation: str | None = Field(default=None, description="问题的标注结果，类似NER格式")
    # 中间数据记录 - 按时间顺序记录所有操作
    operations_history: list[dict] = Field(default_factory=list, description="记录所有操作的历史，包括 get_query_suggestions 和 modify_query")
    # 最终输出结果
    final_output: str | None = Field(default=None, description="Agent的最终输出结果")


# 问题标注 Agent - 三层标注
question_annotation_instructions = """
你是搜索需求分析专家。给定问题（含需求背景），在原文上标注三层：本质、硬性、软性。

## 三层结构

**[本质]** - 问题的核心意图，改变后是完全不同的问题
- 如何获取、教程、推荐、作品、测评等

**[硬]** - 在本质意图下，必须满足的约束
- 地域、时间、对象、质量要求等

**[软]** - 可有可无的修饰
- 能体现、特色、快速、简单等

## 输出格式

词语[本质-描述]、词语[硬-描述]、词语[软-描述]

## 示例

输入：如何获取能体现川西秋季特色的高质量风光摄影素材？
输出：如何获取[本质-找方法] 川西[硬-地域] 秋季[硬-季节] 高质量[硬-质量] 风光摄影素材[硬-对象] 能体现[软-修饰] 特色[软-修饰]

输入：PS抠图教程
输出：PS[硬-工具] 抠图[硬-需求] 教程[本质-学习]

输入：川西秋季风光摄影作品
输出：川西[硬-地域] 秋季[硬-季节] 风光摄影[硬-对象] 作品[本质-欣赏]

## 注意
- 只输出标注后的字符串
- 结合需求背景判断意图
""".strip()

question_annotator = Agent[None](
    name="问题标注专家",
    instructions=question_annotation_instructions,
)

eval_instructions = """
你是搜索query评估专家。给定原始问题标注（三层）和推荐query，评估三个分数。

## 评估目标

用这个推荐query搜索，能否找到满足原始需求的内容？

## 三层评分

### 1. essence_score（本质/意图）= 0 或 1

推荐query的本质/意图是否与原问题一致？

**原问题标注中的[本质-XXX]：**
- 找方法/如何获取 → 推荐词应该是方法/获取途径
- 教程/学习 → 推荐词应该是教程/教学
- 作品/欣赏 → 推荐词应该是作品展示
- 工具/推荐 → 推荐词应该是工具推荐

**评分：**
- 1 = 本质一致
- 0 = 本质改变（完全答非所问）

**示例：**
- 原问题：如何获取[本质-找方法]...素材
- 推荐词：素材获取方法 → essence=1
- 推荐词：素材推荐 → essence=1（都是获取途径）
- 推荐词：素材作品 → essence=0（找方法≠看作品）

### 2. hard_score（硬性约束）= 0 或 1

在本质一致的前提下，是否满足所有硬性约束？

**原问题标注中的[硬-XXX]：**地域、时间、对象、质量、工具（如果用户明确要求）等

**评分：**
- 1 = 所有硬性约束都满足
- 0 = 任一硬性约束不满足

**示例：**
- 原问题：川西[硬-地域] 秋季[硬-季节] 高质量[硬-质量] 风光摄影[硬-对象]
- 推荐词：川西秋季风光摄影 → hard=1
- 推荐词：四川秋季风光摄影 → hard=0（地域泛化：川西→四川）
- 推荐词：川西风光摄影 → hard=0（丢失季节）

### 3. soft_score（软性修饰）= 0-1

软性修饰词保留了多少？

**原问题标注中的[软-XXX]：**修饰词、非关键限定等

**评分参考：**
- 1.0 = 完整保留
- 0.7-0.9 = 保留核心
- 0.4-0.6 = 部分丢失
- 0-0.3 = 大量丢失

## 示例

**原问题标注：** 如何获取[本质-找方法] 川西[硬-地域] 秋季[硬-季节] 高质量[硬-质量] 风光摄影素材[硬-对象] 能体现[软-修饰] 特色[软-修饰]

**推荐query1：** 川西秋季风光摄影素材视频
- essence_score=0（找方法→找素材本身，本质变了）
- hard_score=1（地域、季节、对象都符合）
- soft_score=0.5（丢失"高质量"）
- reason: 本质改变，用户要找获取素材的方法，推荐词是素材内容本身

**推荐query2：** 川西秋季风光摄影素材网站推荐
- essence_score=1（找方法→推荐网站，本质一致）
- hard_score=1（所有硬性约束满足）
- soft_score=0.8（保留核心，"高质量"未明确但推荐通常筛选过）
- reason: 本质一致，硬性约束满足，软性略有丢失但可接受

## 注意

- essence=0 直接拒绝，不管hard/soft多高
- essence=1, hard=0 也要拒绝
- essence=1, hard=1 才看soft_score
""".strip()

class EvaluationFeedback(BaseModel):
    """评估反馈模型 - 三层评分"""
    essence_score: Literal[0, 1] = Field(..., description="本质/意图匹配度，0或1。1=问题本质/意图一致，0=本质改变")
    hard_score: Literal[0, 1] = Field(..., description="硬性约束匹配度，0或1。1=所有硬性约束都满足，0=任一硬性约束不满足")
    soft_score: float = Field(..., description="软性修饰完整度，0-1的浮点数。1.0=完整保留，0.7-0.9=保留核心，0.4-0.6=泛化较大，0-0.3=大量丢失")
    reason: str = Field(..., description="评估理由，包括：1）本质/意图是否一致 2）硬性约束是否满足 3）软性修饰保留情况 4）搜索预期")

evaluator = Agent[None](
    name="评估专家",
    instructions=eval_instructions,
    output_type=EvaluationFeedback,
)

@function_tool
async def get_query_suggestions(wrapper: RunContextWrapper[RunContext], query: str):
    """Fetch search recommendations from Xiaohongshu."""

    # 1. 首次调用时，先标注问题（带需求背景）
    if wrapper.context.question_annotation is None:
        print("正在标注问题...")
        annotation_result = await Runner.run(question_annotator, wrapper.context.q_with_context)
        wrapper.context.question_annotation = str(annotation_result.final_output)
        print(f"问题标注完成：{wrapper.context.question_annotation}")

    # 2. 获取推荐词
    xiaohongshu_api = XiaohongshuSearchRecommendations()
    query_suggestions = xiaohongshu_api.get_recommendations(keyword=query)
    print(f"获取到 {len(query_suggestions) if query_suggestions else 0} 个推荐词：{query_suggestions}")

    # 3. 评估推荐词（三层评分）
    async def evaluate_single_query(q_sug: str):
        """Evaluate a single query suggestion."""
        eval_input = f"""
<原始问题标注（三层）>
{wrapper.context.question_annotation}
</原始问题标注（三层）>

<待评估的推荐query>
{q_sug}
</待评估的推荐query>

请评估该推荐query的三个分数：
1. essence_score: 本质/意图是否一致（0或1）
2. hard_score: 硬性约束是否满足（0或1）
3. soft_score: 软性修饰保留程度（0-1）
4. reason: 详细的评估理由
"""
        evaluator_result = await Runner.run(evaluator, eval_input)
        result: EvaluationFeedback = evaluator_result.final_output
        return {
            "query": q_sug,
            "essence_score": result.essence_score,
            "hard_score": result.hard_score,
            "soft_score": result.soft_score,
            "reason": result.reason,
        }

    # 并发执行所有评估任务
    res = []
    if query_suggestions:
        res = await asyncio.gather(*[evaluate_single_query(q_sug) for q_sug in query_suggestions])
    else:
        res = '未返回任何推荐词'

    # 记录到 RunContext
    wrapper.context.operations_history.append({
        "operation_type": "get_query_suggestions",
        "timestamp": datetime.now().isoformat(),
        "query": query,
        "suggestions": query_suggestions,
        "evaluations": res,
    })

    return res

@function_tool
def modify_query(wrapper: RunContextWrapper[RunContext], 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,
    }

    # 记录到 RunContext
    wrapper.context.operations_history.append({
        "operation_type": "modify_query",
        "timestamp": datetime.now().isoformat(),
        "modification_type": operation_type,
        "original_query": original_query,
        "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."
    }

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

## 核心任务
给定原始问题，通过迭代调用搜索推荐接口（get_query_suggestions），找到满足硬性要求且尽量保留软性信息的推荐query。

## 重要说明
- **你不需要自己评估query的适配性**
- get_query_suggestions 函数会：
  1. 首次调用时自动标注问题（三层：本质、硬性、软性）
  2. 对每个推荐词进行三维度评估
- 返回结果包含：
  - **query**：推荐词
  - **essence_score**：本质/意图匹配度（0或1），0=本质改变，1=本质一致
  - **hard_score**：硬性约束匹配度（0或1），0=不满足约束，1=满足所有约束
  - **soft_score**：软性修饰完整度（0-1），越高表示保留的信息越完整
  - **reason**：详细的评估理由
- **你的职责是分析评估结果，做出决策和策略调整**

## 防止幻觉 - 关键原则
- **严禁编造数据**：只能基于 get_query_suggestions 实际返回的结果进行分析
- **空结果处理**：如果返回的列表为空（[]），必须明确说明"未返回任何推荐词"
- **不要猜测**：在 modify_query 的 reason 中，不能引用不存在的推荐词或评分
- **如实记录**：每次分析都要如实反映实际返回的数据

## 工作流程

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

### 2. 动态探索策略

**第一轮尝试：**
- 使用原始问题直接调用 get_query_suggestions(query="原始问题")
- 第一次调用会自动标注问题（三层），后续调用会复用该标注
- **检查返回结果**：
  - 如果返回空列表 []：说明"该query未返回任何推荐词"，需要简化或替换query
  - 如果有推荐词：查看每个推荐词的 essence_score、hard_score、soft_score 和 reason
- **做出判断**：是否有 essence_score=1 且 hard_score=1 且 soft_score >= 0.7 的推荐词？

**后续迭代：**
如果没有合格推荐词（或返回空列表），必须先调用 modify_query 记录修改，然后再次搜索：

**工具使用流程：**
1. **分析评估反馈**（必须基于实际返回的数据）：
   - **情况A - 返回空列表**：
     * 在 reason 中说明："第X轮未返回任何推荐词，可能是query过于复杂或生僻"
     * 不能编造任何推荐词或评分

   - **情况B - 有推荐词但无合格词**：
     * **首先检查 essence_score**：
       - 如果全是 essence_score=0：本质/意图完全不对，需要重新理解问题
       - 如果有 essence_score=1：本质对了，继续分析

     * **分析 essence_score=1 且 hard_score=1 的推荐词**：
       - 有哪些？soft_score 是多少？
       - 如果 soft_score 较低（<0.7），reason 中说明丢失了哪些信息？
       - 能否通过修改query提高 soft_score？

     * **如果 essence_score=1 但全是 hard_score=0**：
       - reason 中说明了哪些硬性约束不满足？（地域、时间、对象、质量等）
       - 需要如何调整query才能满足硬性约束？

2. **决策修改策略**：基于实际评估反馈，调用 modify_query(original_query, operation_type, new_query, reason)
   - reason 必须引用具体的 essence_score、hard_score、soft_score 和评估理由
   - 不能编造任何数据

3. 使用返回的 new_query 调用 get_query_suggestions

4. 分析新的评估结果，如果仍不满足，重复步骤1-3

**四种操作类型（operation_type）：**
- **简化**：删除冗余词汇，提取核心关键词（当推荐词过于发散时）
- **扩展**：添加限定词或场景描述（当推荐词过于泛化时）
- **替换**：使用同义词、行业术语或口语化表达（当推荐词偏离核心时）
- **组合**：调整关键词顺序或组合方式（当推荐词结构不合理时）

**每次修改的reason必须包含：**
- 上一轮评估结果的关键发现（引用具体的 essence_score、hard_score、soft_score 和评估理由）
- 基于评估反馈，为什么这样修改
- 预期这次修改会带来什么改进

### 3. 决策标准

采用**三级评分标准**：

**优先级1：本质/意图（最高优先级）**
- **essence_score = 1**：本质一致，继续检查
- **essence_score = 0**：本质改变，**直接放弃**

**优先级2：硬性约束（必须满足）**
- **hard_score = 1**：所有约束满足，继续检查
- **hard_score = 0**：约束不满足，**直接放弃**

**优先级3：软性修饰（越高越好）**
- **soft_score >= 0.7**：信息保留较完整，**理想结果**
- **0.5 <= soft_score < 0.7**：有所丢失但可接受，**备选结果**
- **soft_score < 0.5**：丢失过多，继续优化

**采纳标准：**
- **最优**：essence=1 且 hard=1 且 soft >= 0.7
- **可接受**：essence=1 且 hard=1 且 soft >= 0.5（多次尝试后）
- **不可接受**：essence=0 或 hard=0（无论soft多高）

### 4. 迭代终止条件
- **成功终止**：找到 essence=1 且 hard=1 且 soft >= 0.7 的推荐query
- **可接受终止**：5轮后找到 essence=1 且 hard=1 且 soft >= 0.5 的推荐query
- **失败终止**：最多5轮
- **无推荐词**：返回空列表或错误

### 5. 输出要求

**成功找到合格query时：**
```
原始问题：[原问题]
优化后的query：[最终推荐query]
本质匹配度：[essence_score] (1=本质一致)
硬性约束匹配度：[hard_score] (1=所有约束满足)
软性修饰完整度：[soft_score] (0-1)
评估理由：[简要说明]
```

**未找到合格query时：**
```
原始问题：[原问题]
结果：未找到合格推荐query
原因：[本质不符 / 硬性约束不满足 / 软性信息丢失过多]
最接近的推荐词：[如果有essence=1且hard=1的词，列出soft最高的]
建议：[简要建议]
```

## 注意事项
- **第一轮必须使用原始问题**：直接调用 get_query_suggestions(query="原始问题")
  - 第一次调用会自动标注问题（三层），打印出标注结果
  - 后续调用会复用该标注进行评估

- **后续修改必须调用 modify_query**：不能直接用新query调用 get_query_suggestions

- **重点关注评估结果**：每次都要仔细分析返回的三个分数
  - **essence_score=0 直接放弃**，本质不对
  - **hard_score=0 也直接放弃**，约束不满足
  - **优先关注 essence=1 且 hard=1 的推荐词**，分析如何提升 soft_score

- **基于数据决策**：修改策略必须基于评估反馈
  - 引用具体的 essence_score、hard_score、soft_score
  - 引用 reason 中的关键发现

- **采纳标准明确**：
  - **最优**：essence=1 且 hard=1 且 soft >= 0.7，立即采纳
  - **可接受**：essence=1 且 hard=1 且 soft >= 0.5，多次尝试后可采纳
  - **不可接受**：essence=0 或 hard=0，无论soft多高都不能用

- **严禁编造数据**：
  * 如果返回空列表，必须明确说明"未返回任何推荐词"
  * 不能引用不存在的推荐词、分数或评估理由
  * 每次 modify_query 的 reason 必须基于上一轮实际返回的结果
""".strip()




async def main(input_dir: str):
    current_time, log_url = set_trace()

    # 从目录中读取固定文件名
    input_context_file = os.path.join(input_dir, 'context.md')
    input_q_file = os.path.join(input_dir, 'q.md')

    q_context = read_file_as_string(input_context_file)
    q = read_file_as_string(input_q_file)
    q_with_context = f"""
<需求上下文>
{q_context}
</需求上下文>
<当前问题>
{q}
</当前问题>
""".strip()

    # 获取当前文件名作为版本
    version = os.path.basename(__file__)
    version_name = os.path.splitext(version)[0]  # 去掉 .py 后缀

    # 日志保存到输入目录的 output/版本/时间戳 目录下
    log_dir = os.path.join(input_dir, "output", version_name, current_time)

    run_context = RunContext(
        version=version,
        input_files={
            "input_dir": input_dir,
            "context_file": input_context_file,
            "q_file": input_q_file,
        },
        q_with_context=q_with_context,
        q_context=q_context,
        q=q,
        log_dir=log_dir,
        log_url=log_url,
    )
    agent = Agent[RunContext](
        name="Query Optimization Agent",
        instructions=instructions,
        tools=[get_query_suggestions, modify_query],
    )
    result = await Runner.run(agent, input=q_with_context,  context = run_context,)
    print(result.final_output)

    # 保存最终输出到 RunContext
    run_context.final_output = str(result.final_output)

    # 保存 RunContext 到 log_dir
    os.makedirs(run_context.log_dir, exist_ok=True)
    context_file_path = os.path.join(run_context.log_dir, "run_context.json")
    with open(context_file_path, "w", encoding="utf-8") as f:
        json.dump(run_context.model_dump(), f, ensure_ascii=False, indent=2)
    print(f"\nRunContext saved to: {context_file_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="搜索query优化工具")
    parser.add_argument(
        "--input-dir",
        type=str,
        default="input/简单扣图",
        help="输入目录路径，默认: input/简单扣图"
    )
    args = parser.parse_args()

    asyncio.run(main(args.input_dir))