"""
逐 case 提取 capabilities (v5版本)

从 case.json 读取，按 index 遍历每个 case，
调用 LLM 提取 capabilities，按 index 原位回填到 case.json

v5 架构特性：
- 使用结构化 inputs/outputs（role, modality, artifact_type 等10个维度）
- action 对象化：{main_action, mechanism}
- Stage 1 输出 apply_to_draft（自然语言），为 Stage 2 内容树映射做准备
"""

import asyncio
import json
from pathlib import Path
from typing import Any, Dict, Optional, List

from examples.process_pipeline.script.llm_helper import call_llm_with_retry


# v5 词库文件路径
SCRIPT_DIR = Path(__file__).resolve().parent
METHOD_VOCAB_PATH = SCRIPT_DIR / "resource" / "method_vocab_v5.json"

# 默认词库（如果文件不存在时使用）
DEFAULT_METHOD_VOCAB = {
    "流程角色": [
        "生成指令", "编辑指令", "约束条件", "参考素材", "控制信号",
        "区域控制", "参数配置", "模型资源", "源素材", "中间产物",
        "成品", "模板", "评估结果"
    ],
    "模态": ["文本", "图片", "视频", "音频", "特征点", "参数", "模型", "向量", "表格"],
    "主动作": [
        "生成", "编辑", "提取", "改写", "合成", "修复", "增强",
        "训练", "评估", "剪辑", "模板化", "排版", "转写", "配音",
        "匹配", "扩展", "导出"
    ],
    "动作方式": [
        "直接生成", "一致性保持", "结构约束", "质量收束", "局部重绘",
        "扩图", "换背景", "提示词反推", "模板化", "多图融合", "清晰化",
        "风格迁移", "常规编辑", "变体生成", "动画化", "镜头延展",
        "换主体", "换装", "擦除", "调色", "前后景融合", "图文合成",
        "音画合成", "分层叠加", "特征提取", "蒙版提取", "关键帧提取",
        "字幕提取", "风格提取", "片段拼接", "节奏压缩", "转场编排",
        "字幕对齐", "音画同步", "降噪", "补帧", "超分", "稳定化",
        "质感增强", "结构抽象", "变量抽象", "版式套用", "格式转换", "压缩导出"
    ],
}


def load_method_vocab() -> Dict[str, List[str]]:
    """从 JSON 文件加载结构化词库（v5）"""
    if METHOD_VOCAB_PATH.exists():
        try:
            with open(METHOD_VOCAB_PATH, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"Warning: Failed to load method_vocab.json: {e}, using default")
    return DEFAULT_METHOD_VOCAB


def load_prompt_template(prompt_name: str) -> str:
    base_dir = Path(__file__).parent.parent
    prompt_path = base_dir / "prompts" / f"{prompt_name}.prompt"
    with open(prompt_path, "r", encoding="utf-8") as f:
        content = f.read()
    if content.startswith("---"):
        parts = content.split("---", 2)
        if len(parts) >= 3:
            content = parts[2]
    content = content.replace("$system$", "").replace("$user$", "")
    return content.strip()


def render_method_vocab_block(vocab: Dict[str, List[str]]) -> str:
    """渲染结构化接口词库说明（v5）"""
    lines = [
        "\n# 结构化接口词库（v5，必须遵守）",
        "只输出结构化 inputs / outputs / action。",
        "- `role/流程角色` 只写接口职责，不写具体内容 what。",
        "- `modality/模态` 只写媒介或数据形态；统一用 `图片`，不要写 `图像`；统一用 `文本`，不要写 `文字`。",
        "- `artifact_type/工件类型` 写该模态下的具体工件，如 `正向提示词`、`蒙版`。",
        "- `action.main_action` 写主动作；`action.mechanism` 写动作内部机制。",
        "- 只有词库确实不够时才新增术语；新增术语也必须抽象、短、可复用。",
        "",
        "当前词库：",
    ]
    for key, values in vocab.items():
        lines.append(f"- {key}：{'、'.join(values)}")
    return "\n".join(lines)


async def extract_capabilities_from_case_item(
    case_item: Dict[str, Any],
    llm_call: Any,
    model: str = "anthropic/claude-sonnet-4-5"
) -> tuple[Optional[List[Dict[str, Any]]], float]:
    """
    从单个 case item 提取 capabilities (v5版本)

    v5 特性：
    - 结构化 inputs/outputs（role, modality, artifact_type 等）
    - action 对象化：{main_action, mechanism}
    - 输出 apply_to_draft（自然语言），为 Stage 2 内容树映射做准备
    """
    images = case_item.get("images", [])

    case_copy = dict(case_item)
    case_copy.pop("images", None)
    case_copy.pop("_raw", None)
    case_copy.pop("workflow", None)
    case_copy.pop("capabilities", None)

    if not case_copy and not images:
        return None, 0.0

    title = case_item.get("title", "")[:20] or "untitled"
    context = json.dumps(case_copy, ensure_ascii=False, indent=2)

    try:
        prompt_template = load_prompt_template("extract_capability")
        # 添加 v5 词库说明
        method_vocab = load_method_vocab()
        vocab_block = render_method_vocab_block(method_vocab)

        if "%context%" in prompt_template:
            prompt = prompt_template.replace("%context%", context)
        else:
            prompt = prompt_template + f"\n\n案例数据：\n```json\n{context}\n```"

        # 如果 prompt 中有 {interface_vocab} 占位符，替换为词库说明
        if "{interface_vocab}" in prompt:
            prompt = prompt.replace("{interface_vocab}", vocab_block)
        elif vocab_block not in prompt:
            # 如果 prompt 中没有词库说明，添加到末尾
            prompt = prompt + "\n" + vocab_block

    except Exception as e:
        print(f"Warning: Failed to load prompt template: {e}, using fallback")
        method_vocab = load_method_vocab()
        vocab_block = render_method_vocab_block(method_vocab)
        prompt = f"""请从以下案例中提取该案例包含的原子能力，以 JSON 格式输出。

# 输出格式（v5）
{{
  "skip": false,
  "skip_reason": "",
  "capabilities": [
    {{
      "inputs": [
        {{
          "role": "生成指令",
          "modality": "文本",
          "artifact_type": "正向提示词",
          "control_target": ["主体", "场景"],
          "target_scope": ["整图"],
          "constraint_strength": "硬约束",
          "source": "原帖文本",
          "lifecycle": "原始输入",
          "description": "用于触发图片生成的完整提示词"
        }}
      ],
      "outputs": [...],
      "action": {{"main_action": "生成", "mechanism": "直接生成"}},
      "body": "具体做法",
      "effects": ["实现 XX 效果"],
      "stage": ["generate"],
      "tools": [],
      "criterion": null,
      "apply_to_draft": {{"实质": ["相关 what"], "形式": ["相关呈现方式"]}},
      "unstructured_what": []
    }}
  ]
}}

{vocab_block}

案例数据：
{context}

请严格按照上述格式输出JSON，不要包含其他内容。"""

    if images:
        image_urls = [img for img in images[:9] if isinstance(img, str) and img.startswith("http")]
        if image_urls:
            content_array = [{"type": "text", "text": prompt}]
            for url in image_urls:
                content_array.append({"type": "image_url", "image_url": {"url": url}})
            messages = [{"role": "user", "content": content_array}]
        else:
            messages = [{"role": "user", "content": prompt}]
    else:
        messages = [{"role": "user", "content": prompt}]

    result_data, cost = await call_llm_with_retry(
        llm_call=llm_call,
        messages=messages,
        model=model,
        temperature=0.1,
        max_tokens=8000,
        max_retries=3,
        schema_name="extract_capability",
        task_name=f"Capability_{title}",
    )

    # Stage 1 格式：{"skip": bool, "skip_reason": str, "capabilities": [...]}
    # 如果 skip=true，返回空数组
    if not result_data:
        return None, cost

    if result_data.get("skip"):
        return None, cost

    capabilities_data = result_data.get("capabilities", [])
    return capabilities_data, cost


async def extract_capability(
    case_file: Path,
    llm_call: Any,
    model: str = "anthropic/claude-sonnet-4-5",
    max_concurrent: int = 3
) -> Dict[str, Any]:
    """
    按 index 遍历 case.json，提取 capabilities
    """
    with open(case_file, "r", encoding="utf-8") as f:
        case_data = json.load(f)

    cases = case_data.get("cases", [])
    print(f"Extracting capabilities from {len(cases)} cases...")

    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_with_semaphore(case_item):
        async with semaphore:
            index = case_item.get("index", 0)
            raw = case_item.get("_raw", {})
            case_id = raw.get("case_id", "unknown")
            title = case_item.get("title", "")

            print(f"  -> [{index}] [{case_id}] extracting capabilities: {title[:60]}")

            capabilities_data, cost = await extract_capabilities_from_case_item(case_item, llm_call, model)

            status = "ok" if capabilities_data else "null"
            count = len(capabilities_data) if capabilities_data else 0
            print(f"  <- [{index}] [{case_id}] capabilities {status} (count={count})")

            result = dict(case_item)
            result["capabilities"] = capabilities_data
            return result, cost

    tasks = [process_with_semaphore(case) for case in cases]
    results_with_costs = await asyncio.gather(*tasks)

    results = [r[0] for r in results_with_costs]
    costs = [r[1] for r in results_with_costs]
    total_cost = sum(costs)

    success_count = sum(1 for r in results if r.get("capabilities"))
    failed_count = len(results) - success_count

    results.sort(key=lambda x: x.get("index", 0))
    case_data["cases"] = results

    case_file.parent.mkdir(parents=True, exist_ok=True)
    with open(case_file, "w", encoding="utf-8") as f:
        json.dump(case_data, f, ensure_ascii=False, indent=2)

    return {
        "total": len(results),
        "success": success_count,
        "failed": failed_count,
        "total_cost": total_cost,
        "output_file": str(case_file),
    }


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python extract_capability.py <output_dir>")
        sys.exit(1)

    output_dir = Path(sys.argv[1])
    case_file = output_dir / "case.json"

    if not case_file.exists():
        print(f"Error: {case_file} not found")
        sys.exit(1)

    print("Please use this module through run_pipeline.py")