"""
Step3: 基于匹配节点生成灵感点

基于 Step1 的 Top1 匹配结果，以匹配到的人设要素作为锚点，
让 Agent 分析可以产生哪些灵感点
"""
import os
import sys
import json
import asyncio
from pathlib import Path

from agents import Agent, Runner, trace
from agents.tracing.create import custom_span
from lib.my_trace import set_trace_smith as set_trace
from lib.client import get_model
from lib.data_loader import load_persona_data, load_inspiration_list, select_inspiration

# 模型配置
MODEL_NAME = "google/gemini-2.5-pro"


# ========== System Prompt ==========

GENERATE_INSPIRATIONS_PROMPT = """
# 任务
你是一个内容创作者，现在要从一个锚点要素出发，模拟灵感产生的完整思维路径。

## 输入说明
- **<人设体系></人设体系>**: 完整的人设系统，包含所有可用的节点
- **<锚点要素></锚点要素>**: 作为起点的人设要素
- **<要素定义></要素定义>**: 该要素的完整定义
- **<要素上下文></要素上下文>**: 该要素的上下文信息

## 思维路径规则

每个灵感的产生都需要完整的思维路径：**[锚点] → [维度] → [节点/联想]**

1. **[锚点]**: 必须从给定的锚点要素出发
2. **[维度]**: 从锚点特征自然产生的思维方向（动态的，根据锚点特点决定）
   - 例如：跨节点组合、具体场景联想、情感延伸、时间维度、夸张变形等
3. **[节点/联想]**:
   - **[节点]**: 人设体系中的其他节点
   - **[联想]**: 自由联想的具体内容
   - 可以有多个，用 + 连接

## 示例

假设锚点是"宠物迷惑行为"：

```
路径: [锚点]宠物迷惑行为 → [维度]跨节点组合 → [节点]职场疲惫
灵感点: 猫咪上班后的呆滞眼神

路径: [锚点]宠物迷惑行为 → [维度]具体场景联想 → [联想]盯着空气
灵感点: 猫咪对着空气打拳

路径: [锚点]宠物迷惑行为 → [维度]跨节点组合 → [节点]职场疲惫 + [节点]谐音梗
灵感点: 猫咪"喵"班摸鱼
```

## 输出格式（严格JSON）

```json
{
  "灵感点列表": [
    {
      "路径": "[锚点]xxx → [维度]xxx → [节点/联想]xxx",
      "灵感点": "具体的灵感点描述"
    },
    {
      "路径": "[锚点]xxx → [维度]xxx → [联想]xxx",
      "灵感点": "具体的灵感点描述"
    }
  ]
}
```

**要求**：
1. 生成 8-15 个灵感点
2. 路径必须严格按照 `[锚点] → [维度] → [节点/联想]` 格式
3. 维度要多样化，探索不同的思考方向
4. 每个灵感点简洁明确（一句话）
5. 灵感点要具体、可执行、符合人设风格
6. 必须严格按照 JSON 格式输出
""".strip()


def create_agent(model_name: str, prompt: str, name: str) -> Agent:
    """创建 Agent

    Args:
        model_name: 模型名称
        prompt: System prompt
        name: Agent 名称

    Returns:
        Agent 实例
    """
    agent = Agent(
        name=name,
        instructions=prompt,
        model=get_model(model_name),
        tools=[],
    )

    return agent


def parse_json_response(response_content: str, default_value: dict = None) -> dict:
    """解析 JSON 响应

    Args:
        response_content: Agent 返回的响应内容
        default_value: 解析失败时的默认返回值

    Returns:
        解析后的字典
    """
    try:
        # 如果响应包含在 markdown 代码块中，提取 JSON 部分
        if "```json" in response_content:
            json_start = response_content.index("```json") + 7
            json_end = response_content.index("```", json_start)
            json_text = response_content[json_start:json_end].strip()
        elif "```" in response_content:
            json_start = response_content.index("```") + 3
            json_end = response_content.index("```", json_start)
            json_text = response_content[json_start:json_end].strip()
        else:
            json_text = response_content.strip()

        return json.loads(json_text)
    except Exception as e:
        print(f"解析响应失败: {e}")
        return default_value if default_value else {}


def format_persona_system(persona_data: dict) -> str:
    """格式化完整人设系统为文本

    Args:
        persona_data: 人设数据

    Returns:
        格式化的人设系统文本
    """
    lines = ["# 人设系统"]

    # 处理三个部分：灵感点列表、目的点、关键点列表
    for section_key, section_title in [
        ("灵感点列表", "【灵感点】灵感的来源和性质"),
        ("目的点", "【目的点】创作的目的和价值导向"),
        ("关键点列表", "【关键点】内容的核心主体和表达方式")
    ]:
        section_data = persona_data.get(section_key, [])
        if not section_data:
            continue

        lines.append(f"\n## {section_title}\n")

        for perspective in section_data:
            perspective_name = perspective.get("视角名称", "")
            lines.append(f"\n### 视角：{perspective_name}")

            for pattern in perspective.get("模式列表", []):
                pattern_name = pattern.get("分类名称", "")
                pattern_def = pattern.get("核心定义", "")
                lines.append(f"\n  【一级】{pattern_name}")
                if pattern_def:
                    lines.append(f"    定义：{pattern_def}")

                # 二级细分
                for sub in pattern.get("二级细分", []):
                    sub_name = sub.get("分类名称", "")
                    sub_def = sub.get("分类定义", "")
                    lines.append(f"    【二级】{sub_name}：{sub_def}")

    return "\n".join(lines)


def find_element_definition(persona_data: dict, element_name: str) -> str:
    """从人设数据中查找要素的定义

    Args:
        persona_data: 人设数据
        element_name: 要素名称

    Returns:
        要素定义文本，如果未找到则返回空字符串
    """
    # 在灵感点列表中查找
    for section_key in ["灵感点列表", "目的点", "关键点列表"]:
        section_data = persona_data.get(section_key, [])
        for perspective in section_data:
            for pattern in perspective.get("模式列表", []):
                # 检查一级分类
                if pattern.get("分类名称", "") == element_name:
                    definition = pattern.get("核心定义", "")
                    if definition:
                        return definition

                # 检查二级分类
                for sub in pattern.get("二级细分", []):
                    if sub.get("分类名称", "") == element_name:
                        return sub.get("分类定义", "")

    return ""


def find_step1_file(persona_dir: str, inspiration: str, model_name: str) -> str:
    """查找 step1 输出文件

    Args:
        persona_dir: 人设目录
        inspiration: 灵感点名称
        model_name: 模型名称

    Returns:
        step1 文件路径

    Raises:
        SystemExit: 找不到文件时退出
    """
    step1_dir = os.path.join(persona_dir, "how", "灵感点", inspiration)
    model_name_short = model_name.replace("google/", "").replace("/", "_")
    step1_file_pattern = f"*_step1_*_{model_name_short}.json"

    step1_files = list(Path(step1_dir).glob(step1_file_pattern))
    if not step1_files:
        print(f"❌ 找不到 step1 输出文件")
        print(f"查找路径: {step1_dir}/{step1_file_pattern}")
        sys.exit(1)

    return str(step1_files[0])


async def generate_inspirations_with_paths(
    persona_system_text: str,
    matched_element: str,
    element_definition: str,
    element_context: str
) -> list:
    """生成带思维路径的灵感点列表

    Args:
        persona_system_text: 完整人设系统文本
        matched_element: 匹配要素
        element_definition: 要素定义
        element_context: 要素上下文

    Returns:
        灵感点列表 [{"路径": "...", "灵感点": "..."}, ...]
    """
    task_description = f"""## 本次任务

<人设体系>
{persona_system_text}
</人设体系>

<锚点要素>
{matched_element}
</锚点要素>

<要素定义>
{element_definition if element_definition else '无'}
</要素定义>

<要素上下文>
{element_context}
</要素上下文>

请从锚点要素出发，模拟灵感产生的完整思维路径，生成多个灵感点，严格按照 JSON 格式输出。"""

    messages = [{
        "role": "user",
        "content": [{"type": "input_text", "text": task_description}]
    }]

    agent = create_agent(MODEL_NAME, GENERATE_INSPIRATIONS_PROMPT, "Inspiration Path Generator")
    result = await Runner.run(agent, input=messages)

    parsed = parse_json_response(result.final_output, {"灵感点列表": []})
    return parsed.get("灵感点列表", [])


async def process_step3_generate_inspirations(
    step1_top1: dict,
    persona_data: dict,
    current_time: str = None,
    log_url: str = None
) -> dict:
    """执行灵感生成分析（核心业务逻辑 - 单层路径生成）

    Args:
        step1_top1: step1 的 top1 匹配结果
        persona_data: 完整的人设数据
        current_time: 当前时间戳
        log_url: trace URL

    Returns:
        生成结果字典
    """
    # 从 step1 结果中提取信息
    business_info = step1_top1.get("业务信息", {})
    input_info = step1_top1.get("输入信息", {})

    matched_element = business_info.get("匹配要素", "")
    element_context = input_info.get("A_Context", "")

    # 格式化人设系统
    persona_system_text = format_persona_system(persona_data)

    # 查找要素定义
    element_definition = find_element_definition(persona_data, matched_element)

    print(f"\n{'=' * 80}")
    print(f"Step3: 基于锚点生成灵感路径")
    print(f"{'=' * 80}")
    print(f"锚点要素: {matched_element}")
    print(f"要素定义: {element_definition if element_definition else '(未找到定义)'}")
    print(f"模型: {MODEL_NAME}\n")

    # 生成灵感点（单层，包含完整思维路径）
    with custom_span(name="生成带路径的灵感点", data={"锚点要素": matched_element}):
        inspirations = await generate_inspirations_with_paths(
            persona_system_text, matched_element, element_definition, element_context
        )

    print(f"\n{'=' * 80}")
    print(f"完成！共生成 {len(inspirations)} 个灵感点")
    print(f"{'=' * 80}\n")

    # 预览前3个
    if inspirations:
        print("预览前3个灵感点:")
        for i, item in enumerate(inspirations[:3], 1):
            print(f"  {i}. 路径: {item.get('路径', '')}")
            print(f"     灵感: {item.get('灵感点', '')}")
        print()

    # 构建输出
    return {
        "元数据": {
            "current_time": current_time,
            "log_url": log_url,
            "model": MODEL_NAME,
            "步骤": "Step3: 基于匹配节点生成灵感点"
        },
        "锚点信息": {
            "人设要素": matched_element,
            "要素定义": element_definition if element_definition else "无",
            "要素上下文": element_context
        },
        "step1_结果": step1_top1,
        "灵感点列表": inspirations
    }


async def main(current_time: str, log_url: str, force: bool = False):
    """主函数

    Args:
        current_time: 当前时间戳
        log_url: 日志链接
        force: 是否强制重新执行（跳过已存在文件检查）
    """
    # 解析命令行参数
    persona_dir = sys.argv[1] if len(sys.argv) > 1 else "data/阿里多多酱/out/人设_1110"
    inspiration_arg = sys.argv[2] if len(sys.argv) > 2 else "0"

    # 第三个参数：force（如果从命令行调用且有该参数，则覆盖函数参数）
    if len(sys.argv) > 3 and sys.argv[3] == "force":
        force = True

    print(f"{'=' * 80}")
    print(f"Step3: 基于匹配节点生成灵感点")
    print(f"{'=' * 80}")
    print(f"人设目录: {persona_dir}")
    print(f"灵感参数: {inspiration_arg}")

    # 加载数据
    persona_data = load_persona_data(persona_dir)
    inspiration_list = load_inspiration_list(persona_dir)

    # 选择灵感
    try:
        inspiration_index = int(inspiration_arg)
        if 0 <= inspiration_index < len(inspiration_list):
            test_inspiration = inspiration_list[inspiration_index]
            print(f"使用灵感[{inspiration_index}]: {test_inspiration}")
        else:
            print(f"❌ 灵感索引超出范围: {inspiration_index}")
            sys.exit(1)
    except ValueError:
        if inspiration_arg in inspiration_list:
            test_inspiration = inspiration_arg
            print(f"使用灵感: {test_inspiration}")
        else:
            print(f"❌ 找不到灵感: {inspiration_arg}")
            sys.exit(1)

    # 查找并加载 step1 结果
    step1_file = find_step1_file(persona_dir, test_inspiration, MODEL_NAME)
    step1_filename = os.path.basename(step1_file)
    step1_basename = os.path.splitext(step1_filename)[0]

    print(f"Step1 输入文件: {step1_file}")

    # 构建输出文件路径
    output_dir = os.path.join(persona_dir, "how", "灵感点", test_inspiration)
    model_name_short = MODEL_NAME.replace("google/", "").replace("/", "_")
    scope_prefix = step1_basename.split("_")[0]
    result_index = 0
    output_filename = f"{scope_prefix}_step3_top{result_index + 1}_生成灵感_{model_name_short}.json"
    output_file = os.path.join(output_dir, output_filename)

    # 检查文件是否已存在
    if not force and os.path.exists(output_file):
        print(f"\n✓ 输出文件已存在，跳过执行: {output_file}")
        print(f"提示: 如需重新执行，请添加 'force' 参数\n")
        return

    with open(step1_file, 'r', encoding='utf-8') as f:
        step1_data = json.load(f)

    actual_inspiration = step1_data.get("灵感", "")
    step1_results = step1_data.get("匹配结果列表", [])

    if not step1_results:
        print("❌ step1 结果为空")
        sys.exit(1)

    print(f"灵感: {actual_inspiration}")

    # 默认处理 top1
    selected_result = step1_results[result_index]
    print(f"处理第 {result_index + 1} 个匹配结果（Top{result_index + 1}）\n")

    # 执行核心业务逻辑
    output = await process_step3_generate_inspirations(
        step1_top1=selected_result,
        persona_data=persona_data,
        current_time=current_time,
        log_url=log_url
    )

    # 在元数据中添加 step1 匹配索引
    output["元数据"]["step1_匹配索引"] = result_index + 1

    # 保存结果
    os.makedirs(output_dir, exist_ok=True)

    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(output, f, ensure_ascii=False, indent=2)

    # 输出统计信息
    stats = output.get("统计", {})
    print(f"\n{'=' * 80}")
    print(f"统计信息:")
    print(f"  外层维度数: {stats.get('外层维度数', 0)}")
    print(f"  内层维度数: {stats.get('内层维度数', 0)}")
    print(f"  灵感点总数: {stats.get('灵感点总数', 0)}")
    print(f"{'=' * 80}")

    print(f"\n完成！结果已保存到: {output_file}")
    if log_url:
        print(f"Trace: {log_url}\n")


if __name__ == "__main__":
    # 设置 trace
    current_time, log_url = set_trace()

    # 使用 trace 上下文包裹整个执行流程
    with trace("Step3: 生成灵感点"):
        asyncio.run(main(current_time, log_url))