"""
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 = """
# 任务
你是一个内容创作者，现在要从一个锚点分类出发，通过思维路径推导出可能触发创作冲动的客观刺激源（灵感点）。

## 核心概念

**分类**（维度）：创作者接收外界信息刺激的角度或通道
- 格式：2-4个字，简洁直观

**灵感点**：创作前遇到的、触发创作冲动的客观刺激源
- 本质：作者被动接收的信息（看到的、听说的、发现的、观察到的、感知到的）
- 格式：不超过15个字，使用自然、通俗、口语化的表达
- 表达要求：
  * 使用日常生活语言，避免学术化、抽象化词汇堆砌
  * 优先使用"的"字短语（如"夏日的热闹景象"）或动宾短语（如"观察到的自然互动"）
  * 禁止使用多个抽象名词连用（如"具象化动态互动自然拟人"）
  * 让普通人一看就懂

**描述**：对刺激源本身是什么的详细说明
- 描述刺激源的具体特征、形态、场景、内容等客观信息
- 注意区分：刺激源内容本身 vs 呈现方式/表现形式

**推理路径**：展示从锚点分类到灵感点的推导过程
- 格式：`[锚点分类] → [思维方向] → [联想节点] → [灵感点]`
- 思维方向：从锚点出发的联想角度（如：具体场景、情感延伸、时间维度、反差对比等）
- 联想节点：人设体系中的相关节点，或具体的联想内容

## 严格禁止
- 不描述创作者如何运用/展现/表达刺激，不使用推理性词汇
- 不能是创作形式、表现手法、表达方式、呈现方式、风格、格式等
- 必须是被动接收的刺激，不能是主动创造的内容
- 不解释创作者为什么被触发、如何使用
- 不进行主观推理和价值判断
- 禁止词汇堆砌

## 输入说明
- **<人设体系></人设体系>**: 完整的人设系统，包含所有可用节点
- **<锚点分类></锚点分类>**: 作为起点的分类维度（接收刺激的角度）
- **<分类定义></分类定义>**: 该分类的完整定义
- **<分类上下文></分类上下文>**: 该分类的上下文信息

## 推导方法

从锚点分类出发，通过思维路径推导灵感点：

1. **确定锚点**：锚点分类是什么？
2. **选择思维方向**：从这个分类可以往哪个方向联想？
3. **找到联想节点**：结合人设体系，这个方向上有哪些相关节点或内容？
4. **得出灵感点**：这些联想最终指向什么具体的客观刺激？

## 输出格式（严格JSON）

**重要：必须输出严格的 JSON 格式，注意以下几点：**
- 使用英文双引号 `"` 而非中文引号 `""`
- 字段值中如果包含引号，必须转义 `\"`
- 不要在最后一个元素后添加逗号
- 确保所有括号正确闭合
- 描述内容不要换行，保持在一行内

```json
{
  "灵感点列表": [
    {
      "推理路径": "[锚点分类] → [思维方向] → [联想节点] → [灵感点]",
      "灵感点": "具体的客观刺激源描述（不超过15字，口语化）",
      "描述": "对这个刺激源本身是什么的详细说明，描述其具体特征、形态、场景、内容等客观信息（不换行，一句话）"
    }
  ]
}
```

**要求**：
1. 生成 8-15 个灵感点
2. 每个灵感点必须是客观刺激源，不能是创作手法
3. "推理路径"字段：清晰展示推导过程
4. "灵感点"字段：简洁口语化，不超过15字
5. "描述"字段：客观描述刺激源本身，不涉及如何运用，不换行
6. 字段值避免使用特殊字符（如未转义的引号、换行符等）
7. 必须输出完整有效的 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:
        解析后的字典
    """
    import re

    # 提取 JSON 文本
    def extract_json_text(content):
        if "```json" in content:
            json_start = content.index("```json") + 7
            # 查找下一个 ``` 或 ``` 后的内容结束
            try:
                json_end = content.index("```", json_start)
            except ValueError:
                # 如果找不到结束标记，取到末尾
                json_end = len(content)
            return content[json_start:json_end].strip()
        elif "```" in content:
            json_start = content.index("```") + 3
            try:
                json_end = content.index("```", json_start)
            except ValueError:
                json_end = len(content)
            return content[json_start:json_end].strip()
        else:
            return content.strip()

    json_text = extract_json_text(response_content)

    # 尝试1: 直接解析
    try:
        return json.loads(json_text)
    except json.JSONDecodeError as e:
        print(f"\n⚠️  JSON 解析失败（尝试1），开始修复...")
        print(f"   错误: {e}\n")

    # 尝试2: 修复常见问题
    try:
        # 修复1: 去除尾部逗号
        fixed = re.sub(r',(\s*[}\]])', r'\1', json_text)

        # 修复2: 处理未完成的JSON（截断问题）
        # 如果JSON被截断了，尝试补全
        if fixed.count('{') > fixed.count('}'):
            # 补充缺失的闭合括号
            diff = fixed.count('{') - fixed.count('}')
            fixed += '\n' + '  }'*diff
        if fixed.count('[') > fixed.count(']'):
            diff = fixed.count('[') - fixed.count(']')
            fixed += '\n' + '  ]'*diff

        # 修复3: 去除未完成的最后一项
        # 如果最后一项没有闭合，移除它
        lines = fixed.split('\n')
        # 倒序查找最后一个完整的对象
        bracket_count = 0
        last_complete_idx = len(lines)
        for i in range(len(lines) - 1, -1, -1):
            line = lines[i]
            bracket_count += line.count('}') - line.count('{')
            bracket_count += line.count(']') - line.count('[')
            if bracket_count == 0 and ('}' in line or ']' in line):
                last_complete_idx = i + 1
                break

        if last_complete_idx < len(lines):
            print(f"   检测到未完成的内容，截断到第 {last_complete_idx} 行")
            fixed = '\n'.join(lines[:last_complete_idx])

        result = json.loads(fixed)
        print(f"✓ JSON 修复成功\n")
        return result
    except Exception as fix_error:
        print(f"   修复失败: {fix_error}\n")

    # 最终失败，返回默认值
    print(f"\n{'!' * 80}")
    print(f"⚠️  所有尝试均失败，返回空结果")
    print(f"{'!' * 80}")
    print(f"\n原始响应内容:\n")
    print(response_content[:3000])
    print(f"\n{'!' * 80}\n")
    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,
    anchor_category: str,
    category_definition: str,
    category_context: str
) -> list:
    """从锚点分类推导灵感点列表

    Args:
        persona_system_text: 完整人设系统文本
        anchor_category: 锚点分类（维度）
        category_definition: 分类定义
        category_context: 分类上下文

    Returns:
        灵感点列表 [{"分类": "...", "灵感点": "...", "描述": "...", "推理": "..."}, ...]
    """
    task_description = f"""## 本次任务

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

<锚点分类>
{anchor_category}
</锚点分类>

<分类定义>
{category_definition if category_definition else '无'}
</分类定义>

<分类上下文>
{category_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("输入信息", {})

    anchor_category = business_info.get("匹配要素名称", "")
    category_context = input_info.get("A_Context", "")

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

    # 查找分类定义
    category_definition = find_element_definition(persona_data, anchor_category)

    print(f"\n{'=' * 80}")
    print(f"Step3: 从锚点分类推导灵感点")
    print(f"{'=' * 80}")
    print(f"锚点分类: {anchor_category}")
    print(f"分类定义: {category_definition if category_definition else '(未找到定义)'}")
    print(f"模型: {MODEL_NAME}\n")

    # 生成灵感点
    with custom_span(name="从锚点分类推导灵感点", data={"锚点分类": anchor_category}):
        inspirations = await generate_inspirations_with_paths(
            persona_system_text, anchor_category, category_definition, category_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('灵感点', '')} ({len(item.get('灵感点', ''))}字)")
            print(f"     描述: {item.get('描述', '')[:60]}...")
        print()

    # 构建输出
    return {
        "元数据": {
            "current_time": current_time,
            "log_url": log_url,
            "model": MODEL_NAME,
            "步骤": "Step3: 从锚点分类推导灵感点"
        },
        "锚点信息": {
            "锚点分类": anchor_category,
            "分类定义": category_definition if category_definition else "无",
            "分类上下文": category_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)

    # 输出统计信息
    inspirations_list = output.get("灵感点列表", [])
    print(f"\n{'=' * 80}")
    print(f"统计信息:")
    print(f"  生成灵感点数量: {len(inspirations_list)}")

    # 统计字段完整性
    complete_count = sum(
        1 for item in inspirations_list
        if all(key in item and item[key] for key in ["推理路径", "灵感点", "描述"])
    )
    print(f"  字段完整的灵感点: {complete_count}/{len(inspirations_list)}")

    # 统计灵感点字数
    lengths = [len(item.get("灵感点", "")) for item in inspirations_list if item.get("灵感点")]
    if lengths:
        avg_length = sum(lengths) / len(lengths)
        max_length = max(lengths)
        over_15 = sum(1 for l in lengths if l > 15)
        print(f"  灵感点字数: 平均 {avg_length:.1f}字, 最长 {max_length}字")
        if over_15 > 0:
            print(f"  ⚠️  超过15字的灵感点: {over_15}个")
    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))