# 3-Stage IR Generation Pipeline
#
# Stage 1: Semantic Index — full document → function_units + concepts
# Stage 2: Per-Unit IR Extraction — precision context → detailed IR rules
# Stage 3: Deterministic Merge & Audit — dedup, assign IDs, coverage report

import concurrent.futures
import hashlib
import json
import logging
import re
import time
import uuid
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

from server.config import settings
from server.core.llm_provider.router import router

logger = logging.getLogger(__name__)

PROMPTS_DIR = Path(__file__).parent.parent / "prompts"  # services/prompts/
MAX_CONCURRENCY = 5

PASS, WARN, FAIL = "[PASS]", "[WARN]", "[FAIL]"


# =============================================================================
# Stage 1: Semantic Index Generation
# =============================================================================

def format_document_for_prompt(doc: dict) -> str:
    """Render the full parsed document as a readable string for the LLM prompt."""
    lines = ["=== SECTIONS ==="]

    for i, section in enumerate(doc.get("sections", [])):
        source = section.get("source", f"(无标题-章节{i})")
        lines.append(f"\n--- Section: {source} ---")

        for block in section.get("blocks", []):
            if block["type"] == "para":
                lines.append(f"[段落 {block['index']}] {block['text']}")
            elif block["type"] == "table":
                lines.append(f"[表格 {block.get('table', '?')}]")
                headers = block.get("headers", [])
                lines.append(f"  表头: {' | '.join(headers)}")
                for row in block.get("rows", []):
                    cols = row.get("columns", [])
                    cell_texts = []
                    for c in cols:
                        cell_texts.append(f"[行{c.get('row','?')}]{c.get('name','')}: {c.get('text','')}")
                    lines.append(f"  {'; '.join(cell_texts)}")

        images = section.get("images", [])
        if images:
            lines.append(f"  图片引用: {', '.join(images)}")

    # Image Analysis
    lines.append("\n\n=== IMAGE_ANALYSIS ===")
    for img in doc.get("image_analysis", []):
        rid = img.get("rid", "?")
        img_type = img.get("type", "?")
        lines.append(f"\n--- Image: {rid} (type={img_type}) ---")
        lines.append(f"  描述: {img.get('description', '')[:500]}")

    # Resolved Conflicts
    conflicts = doc.get("resolved_conflicts", [])
    if conflicts:
        lines.append("\n\n=== RESOLVED_CONFLICTS ===")
        for c in conflicts:
            lines.append(
                f"  [{c.get('conflict_type','?')}] {c.get('section','?')}: "
                f"以{c.get('source','?')}为准 — {c.get('correction','')}"
            )

    return "\n".join(lines)


def extract_json_from_response(text: str, expect_array: bool = False) -> str:
    """Robustly extract JSON object or array from LLM response."""
    open_char = "[" if expect_array else "{"
    close_char = "]" if expect_array else "}"

    # Try code fence first
    pattern = rf"```(?:json)?\s*(\[{open_char}[\s\S]*?\{close_char}\])\s*```" if expect_array else rf"```(?:json)?\s*(\{{[\s\S]*?\}})\s*```"
    m = re.search(pattern, text)
    if m:
        return m.group(1).strip()

    # Find outermost bracket
    start = text.find(open_char)
    if start == -1:
        raise ValueError(f"No JSON {'array' if expect_array else 'object'} found in LLM response")

    depth = 0
    for i in range(start, len(text)):
        if text[i] == open_char:
            depth += 1
        elif text[i] == close_char:
            depth -= 1
            if depth == 0:
                return text[start: i + 1]

    raise ValueError(f"Unclosed JSON {'array' if expect_array else 'object'} in LLM response")


async def stage1_semantic_index(doc: dict) -> dict:
    """Stage 1: Generate semantic index from full document."""
    logger.info("[PIPELINE] ═══ Stage 1/3: 语义索引 ═══")

    template_path = PROMPTS_DIR / "ir_step1_semantic_index.txt"
    if not template_path.exists():
        logger.warning("[PIPELINE] Stage 1 prompt template not found, using fallback")
        return _fallback_semantic_index(doc)

    template = template_path.read_text(encoding="utf-8")
    formatted_doc = format_document_for_prompt(doc)
    prompt = template.replace("{document_json}", formatted_doc)

    logger.info("[PIPELINE] Prompt: %d chars (~%d tokens)", len(prompt), len(prompt) // 3)

    client = router.get_text_client()
    for attempt in range(3):
        try:
            raw = client.chat(
                model=router.text_model,
                messages=[
                    {"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON，不输出任何其他文字。"},
                    {"role": "user", "content": prompt},
                ],
                temperature=0.1,
            )
            json_str = extract_json_from_response(raw)
            result = json.loads(json_str)

            # Validate
            if "function_units" not in result or not result["function_units"]:
                raise ValueError("Missing function_units")
            logger.info("[PIPELINE] ✅ Stage 1 完成: %d function_units, %d concepts",
                         len(result["function_units"]), len(result.get("concepts", [])))
            return result

        except Exception as e:
            logger.warning("[PIPELINE] Stage 1 attempt %d failed: %s", attempt + 1, e)
            if attempt < 2:
                time.sleep(2)

    # Fallback
    logger.warning("[PIPELINE] Stage 1 全部重试失败 exhausted, using fallback")
    return _fallback_semantic_index(doc)


def _fallback_semantic_index(doc: dict) -> dict:
    """Generate a basic semantic index without LLM (for mock/demo)."""
    units = []
    unit_idx = 1

    for section in doc.get("sections", []):
        source = section.get("source", "")
        for block in section.get("blocks", []):
            if block["type"] == "para":
                # Create a function unit from each substantial paragraph
                text = block.get("text", "")
                if len(text) > 30:
                    units.append({
                        "unit_id": f"FU-{unit_idx:03d}",
                        "name": text[:60],
                        "description": text[:200],
                        "sources": [{"section": source, "type": "para", "text_snippet": text[:200]}],
                    })
                    unit_idx += 1
            elif block["type"] == "table":
                for row in block.get("rows", []):
                    cols = row.get("columns", [])
                    texts = [f"{c.get('name','')}: {c.get('text','')}" for c in cols if c.get('text')]
                    if texts:
                        combined = "; ".join(texts)
                        units.append({
                            "unit_id": f"FU-{unit_idx:03d}",
                            "name": combined[:60],
                            "description": combined[:200],
                            "sources": [{"section": source, "type": "table", "row": row.get("columns", [{}])[0].get("row", 1) if cols else 1, "text_snippet": combined[:200]}],
                        })
                        unit_idx += 1

    return {
        "feature_name": doc.get("sections", [{}])[0].get("source", "未命名功能"),
        "concepts": [],
        "function_units": units,
    }


# =============================================================================
# Stage 2: Per-Unit IR Extraction
# =============================================================================

def build_context_package(fu: dict, doc: dict, image_sources: dict) -> dict:
    """Build a precision context package for one function unit."""
    texts = []
    tables = []
    images = []
    seen_sections = set()

    img_by_rid = {img.get("rid", ""): img for img in doc.get("image_analysis", [])}

    for src in fu.get("sources", []):
        section_key = src.get("section", "")
        src_type = src.get("type", "")

        if src_type in ("para", "table") and section_key:
            if section_key in seen_sections:
                continue
            seen_sections.add(section_key)

            # Find matching section
            for section in doc.get("sections", []):
                if section.get("source", "") == section_key:
                    for block in section.get("blocks", []):
                        if block["type"] == "para":
                            texts.append({
                                "section": section_key,
                                "text": block["text"],
                            })
                        elif block["type"] == "table":
                            all_rows = []
                            for row in block.get("rows", []):
                                for col in row.get("columns", []):
                                    all_rows.append({
                                        "row": col.get("row"),
                                        "name": col.get("name"),
                                        "text": col.get("text"),
                                    })
                            # Trim: max 20 cells per table to reduce prompt size
                            tables.append({
                                "section": section_key,
                                "headers": block.get("headers", []),
                                "all_rows": all_rows[:20],
                                "_truncated": len(all_rows) > 20,
                            })
                    break

        if src_type in ("image", "logic_tree"):
            img_id = src.get("image_id", "")
            if img_id and img_id in img_by_rid:
                img = img_by_rid[img_id]
                images.append({
                    "image_id": img_id,
                    "type": img.get("type", "other"),
                    "description": img.get("description", ""),
                })

    return {
        "unit_id": fu["unit_id"],
        "unit_name": fu.get("name", ""),
        "unit_description": fu.get("description", ""),
        "texts": texts,
        "tables": tables,
        "images": images,
        "resolved_conflicts": doc.get("resolved_conflicts", []),
    }


async def stage2_extract_rules(semantic_index: dict, doc: dict) -> list[dict]:
    """Stage 2: Extract detailed IR rules for each function unit."""
    logger.info("=" * 50)
    logger.info("[PIPELINE] ═══ Stage 2/3: Per-Unit IR Extraction")

    template_path = PROMPTS_DIR / "ir_step2_extraction.txt"
    if not template_path.exists():
        return _fallback_rules(semantic_index)

    template = template_path.read_text(encoding="utf-8")
    function_units = semantic_index.get("function_units", [])
    image_sources = doc.get("image_sources", {})

    logger.info("Processing %d function units (max concurrency: %d)", len(function_units), MAX_CONCURRENCY)

    # Build context packages
    packages = []
    for fu in function_units:
        pkg = build_context_package(fu, doc, image_sources)
        packages.append(pkg)

    # Parallel LLM calls
    fragments = []
    client = router.get_text_client()

    def extract_one(pkg: dict) -> dict:
        """Extract rules for one function unit."""
        prompt = _build_stage2_prompt(template, pkg)
        unit_id = pkg["unit_id"]

        for attempt in range(3):
            try:
                raw = client.chat(
                    model=router.text_model,
                    messages=[
                        {"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON 数组，不输出任何其他文字。"},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.1,
                )
                json_str = extract_json_from_response(raw, expect_array=True)
                rules = json.loads(json_str)
                if not isinstance(rules, list):
                    raise ValueError(f"Expected array, got {type(rules).__name__}")
                return {"unit_id": unit_id, "unit_name": pkg["unit_name"], "rules": rules}
            except Exception as e:
                if attempt < 2:
                    time.sleep(2)
                else:
                    logger.warning("Unit %s failed after 3 attempts: %s", unit_id, e)
                    return {"unit_id": unit_id, "unit_name": pkg["unit_name"], "rules": [], "error": str(e)}
        return {"unit_id": unit_id, "unit_name": pkg["unit_name"], "rules": [], "error": "unknown"}

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_CONCURRENCY) as executor:
        futures = {executor.submit(extract_one, pkg): pkg for pkg in packages}
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            fragments.append(result)
            n_rules = len(result.get("rules", []))
            logger.info("[PIPELINE]   %s: %d rules %s", result["unit_id"], n_rules,
                         "❌" if result.get("error") else "✅")

    fragments.sort(key=lambda f: f["unit_id"])
    total = sum(len(f.get("rules", [])) for f in fragments)
    logger.info("[PIPELINE] ✅ Stage 2 完成: %d fragments, %d total rules", len(fragments), total)
    return fragments


def _build_stage2_prompt(template: str, pkg: dict) -> str:
    """Render the stage 2 prompt template with context package data."""
    # Simple str.replace (avoids format issues with curly braces in JSON)
    prompt = template
    prompt = prompt.replace("{unit_id}", pkg["unit_id"])
    prompt = prompt.replace("{unit_name}", pkg["unit_name"])
    prompt = prompt.replace("{unit_description}", pkg["unit_description"])
    prompt = prompt.replace("{texts}", json.dumps(pkg.get("texts", []), ensure_ascii=False, indent=2))
    prompt = prompt.replace("{tables}", json.dumps(pkg.get("tables", []), ensure_ascii=False, indent=2))
    prompt = prompt.replace("{images}", json.dumps(pkg.get("images", []), ensure_ascii=False, indent=2))
    prompt = prompt.replace("{resolved_conflicts}", json.dumps(pkg.get("resolved_conflicts", []), ensure_ascii=False, indent=2))
    return prompt


def _fallback_rules(semantic_index: dict) -> list[dict]:
    """Generate basic rules without LLM."""
    fragments = []
    for fu in semantic_index.get("function_units", []):
        fragments.append({
            "unit_id": fu["unit_id"],
            "unit_name": fu.get("name", ""),
            "rules": [{
                "description": fu.get("description", ""),
                "priority": "P2",
                "sources": fu.get("sources", []),
                "precondition": {},
                "trigger": {"operator": "AND", "conditions": []},
                "actions": [],
            }],
        })
    return fragments


# =============================================================================
# Stage 3: Merge & Audit
# =============================================================================

def rule_signature(rule: dict) -> str:
    """SHA256 signature from trigger + actions for dedup."""
    trigger = rule.get("trigger", {})
    actions = rule.get("actions", [])
    conditions = sorted(trigger.get("conditions", []), key=lambda c: json.dumps(c, sort_keys=True))
    sorted_actions = sorted(actions, key=lambda a: json.dumps(a, sort_keys=True))
    sig = json.dumps({"conditions": conditions, "actions": sorted_actions}, ensure_ascii=False, sort_keys=True)
    return hashlib.sha256(sig.encode()).hexdigest()[:16]


def stage3_merge_and_audit(fragments: list[dict], semantic_index: dict, doc: dict) -> dict:
    """Stage 3: Merge rules, assign IDs, and generate audit report."""
    logger.info("=" * 50)
    logger.info("[PIPELINE] ═══ Stage 3/3: Merge & Audit")

    # Merge and dedup
    sig_map: dict[str, dict] = {}
    order = []
    for f in fragments:
        for rule in f.get("rules", []):
            sig = rule_signature(rule)
            if sig in sig_map:
                existing = sig_map[sig]
                existing_sources = existing.setdefault("sources", [])
                for src in rule.get("sources", []):
                    if src not in existing_sources:
                        existing_sources.append(src)
                if len(rule.get("description", "")) > len(existing.get("description", "")):
                    existing["description"] = rule["description"]
            else:
                sig_map[sig] = dict(rule)
                order.append(sig)

    merged = [sig_map[sig] for sig in order]
    logger.info("[PIPELINE] 合并: %d → %d rules (dedup)", sum(len(f.get("rules", [])) for f in fragments), len(merged))

    # Assign rule IDs: IR-{feature_id}-{section}-{seq}
    section_counters = defaultdict(int)
    feature_name = semantic_index.get("feature_name", "FEATURE")
    feature_id = re.sub(r'[^A-Z]', '', feature_name.upper())[:6] or "FEAT"
    if len(feature_id) < 2:
        feature_id = (feature_id + "FEATURE")[:4]

    for rule in merged:
        # Extract section from first para/table source, fallback to "0"
        section = "0"
        for src in rule.get("sources", []):
            s = src.get("section", "")
            if s:
                # Extract first number pattern like "4.2.1" or "3.1"
                m = re.search(r'[\d]+(?:\.[\d]+)*', s)
                if m:
                    section = m.group()
                    break
        section_counters[section] += 1
        rule["rule_id"] = f"IR-{feature_id}-{section}-{section_counters[section]:03d}"

    # Build final IR
    ir_final = {
        "feature": feature_name,
        "feature_id": feature_id,
        "rules": merged,
        "rule_count": len(merged),
        "generated_at": datetime.now(timezone.utc).isoformat(),
    }

    # Generate audit report
    audit_results = _generate_audit(merged, doc)
    report = _format_audit_report(feature_name, len(merged), audit_results)

    logger.info("[PIPELINE] ✅ Stage 3 完成: %d rules, %d audit issues", len(merged), audit_results.get("issue_count", 0))
    return {"ir": ir_final, "audit": audit_results, "audit_report": report}


def _generate_audit(rules: list[dict], doc: dict) -> dict:
    """Run coverage audits and return structured results."""
    results = []

    # 1. Source coverage: check that all sections with content are referenced
    all_sources = set()
    for rule in rules:
        for src in rule.get("sources", []):
            section = src.get("section", "")
            if section:
                all_sources.add(section)

    doc_sections = set()
    for section in doc.get("sections", []):
        source = section.get("source", "")
        if source and section.get("blocks"):
            doc_sections.add(source)

    uncovered = doc_sections - all_sources
    results.append({
        "check": "文档章节覆盖",
        "status": PASS if not uncovered else WARN,
        "detail": f"已覆盖 {len(all_sources)}/{len(doc_sections)} 章节"
                  + (f"; 未覆盖: {list(uncovered)[:5]}" if uncovered else ""),
    })

    # 2. Image coverage
    img_count = len(doc.get("image_analysis", []))
    img_referenced = set()
    for rule in rules:
        for src in rule.get("sources", []):
            if src.get("type") in ("image", "logic_tree"):
                img_referenced.add(src.get("image_id", ""))
    results.append({
        "check": "图片覆盖率",
        "status": PASS if len(img_referenced) >= img_count else (WARN if img_referenced else FAIL),
        "detail": f"已引用 {len(img_referenced)}/{img_count} 张图片",
    })

    # 3. Priority distribution
    p0 = sum(1 for r in rules if r.get("priority") == "P0")
    p1 = sum(1 for r in rules if r.get("priority") == "P1")
    p2 = sum(1 for r in rules if r.get("priority") == "P2")
    results.append({
        "check": "优先级分布",
        "status": PASS if p0 > 0 else WARN,
        "detail": f"P0={p0}, P1={p1}, P2={p2}",
    })

    issue_count = sum(1 for r in results if r["status"] in (WARN, FAIL))
    return {"checks": results, "issue_count": issue_count}


def _sync_call(async_fn, *args, **kwargs):
    """Run an async function synchronously using the existing event loop if available."""
    import asyncio
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        # No event loop running — just use asyncio.run
        return asyncio.run(async_fn(*args, **kwargs))
    else:
        # Event loop is running (e.g., inside FastAPI) — must use run_until_complete
        # Create a new loop in a thread to avoid conflicts
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            future = ex.submit(asyncio.run, async_fn(*args, **kwargs))
            return future.result()


def _format_audit_report(feature_name: str, rule_count: int, audit: dict) -> str:
    """Generate Markdown audit report."""
    lines = [
        f"# IR 完整性审计报告",
        f"",
        f"**功能**: {feature_name}",
        f"**规则总数**: {rule_count}",
        f"**生成时间**: {datetime.now(timezone.utc).isoformat()}",
        f"",
        f"> ⚠️ 请人工审查以下 ⚠️ 和 ❌ 项。",
        f"",
        f"## 审计结果",
        f"",
        f"| 检查项 | 状态 | 详情 |",
        f"|--------|------|------|",
    ]
    for r in audit.get("checks", []):
        lines.append(f"| {r['check']} | {r['status']} | {r['detail']} |")
    lines.append("")
    lines.append(f"**审计问题总计**: {audit.get('issue_count', 0)} 个需要人工审查")
    return "\n".join(lines)


# =============================================================================
# Pipeline Orchestrator
# =============================================================================

class IRPipeline:
    """Orchestrates the 3-stage IR generation pipeline."""

    def __init__(self):
        self._store: dict[str, dict] = {}

    def run_streaming(self, prd_id: str, parsed_doc: dict):
        """Run pipeline as a sync generator yielding progress events.

        This is a plain generator (NOT async) so it never blocks the event loop
        between yields. The SSE endpoint wraps it in a thread + Queue bridge.
        """
        t_total_start = time.time()
        n_sections = len(parsed_doc.get("sections", []))
        est_units = max(n_sections * 2, 3)
        est_total = 20 + (est_units * 20 / 3) + 1

        # --- Stage 1 ---
        yield {"stage": 1, "status": "running", "message": "Stage 1: 分析文档结构，识别功能单元...",
               "elapsed": 0, "stage_total": 3, "detail": f"文档含 {n_sections} 个章节"}
        t0 = time.time()
        try:
            semantic_index = _sync_call(stage1_semantic_index, parsed_doc)
        except Exception as e:
            yield {"stage": 1, "status": "error", "message": f"Stage 1 失败: {e}", "stage_total": 3}
            return
        t1 = time.time()
        n_units = len(semantic_index.get("function_units", []))
        est_total = 20 + (max(n_units, 1) * 25 / 3) + 1
        yield {"stage": 1, "status": "done", "message": f"Stage 1 完成: 识别 {n_units} 个功能单元",
               "elapsed": round(t1 - t0, 1), "stage_total": 3, "detail": f"识别 {n_units} 个功能单元",
               "n_units": n_units, "estimated_total": round(est_total)}

        # --- Stage 2 ---
        est_stage2 = max(n_units, 1) * 25 / min(MAX_CONCURRENCY, max(n_units, 1))
        yield {"stage": 2, "status": "running", "message": f"Stage 2: 逐单元提取 IR 规则 ({n_units} 个单元)...",
               "elapsed": 0, "stage_total": 3, "detail": f"预计 {est_stage2:.0f}s，{n_units} 个单元 × {min(MAX_CONCURRENCY, n_units)} 并发",
               "n_units": n_units, "estimated_total": round(est_total)}
        t1b = time.time()
        completed = 0
        total_rules = 0

        fragments = []
        function_units = semantic_index.get("function_units", [])

        # Optimization: filter out trivial units (very short desc, no text sources)
        filtered_units = []
        for fu in function_units:
            sources = fu.get("sources", [])
            desc = fu.get("description", "")
            has_text = any(s.get("type") in ("para", "table") for s in sources)
            if len(desc) < 15 and not has_text:
                logger.info("[PIPELINE] 跳过空单元: %s", fu.get("unit_id", "?"))
                continue
            filtered_units.append(fu)

        skipped = len(function_units) - len(filtered_units)
        if skipped:
            logger.info("[PIPELINE] Stage 2 过滤掉 %d 个空单元，剩余 %d", skipped, len(filtered_units))
        function_units = filtered_units

        image_sources = parsed_doc.get("image_sources", {})
        template_path = PROMPTS_DIR / "ir_step2_extraction.txt"
        template = template_path.read_text(encoding="utf-8") if template_path.exists() else ""
        client = router.get_text_client()
        packages = [build_context_package(fu, parsed_doc, image_sources) for fu in function_units]

        with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_CONCURRENCY) as executor:
            futures = {}
            for i, pkg in enumerate(packages):
                future = executor.submit(self._extract_one_sync, template, pkg, client)
                futures[future] = (i, pkg["unit_id"], pkg["unit_name"])

            for future in concurrent.futures.as_completed(futures):
                i, uid, uname = futures[future]
                try:
                    rules = future.result()
                    fragments.append({"unit_id": uid, "unit_name": uname, "rules": rules})
                    total_rules += len(rules)
                except Exception as e:
                    logger.warning("Unit %s failed: %s", uid, e)
                    fragments.append({"unit_id": uid, "unit_name": uname, "rules": [], "error": str(e)})

                completed += 1
                elapsed = round(time.time() - t1b, 1)
                yield {"stage": 2, "status": "running",
                       "message": f"Stage 2: {completed}/{n_units} 单元完成 ({total_rules} 条规则)",
                       "elapsed": elapsed, "stage_total": 3,
                       "detail": f"已完成 {completed}/{n_units}，累计 {total_rules} 条规则",
                       "completed": completed, "total": n_units, "rules_so_far": total_rules,
                       "estimated_total": round(est_total)}

        fragments.sort(key=lambda f: f["unit_id"])
        t2 = time.time()
        logger.info("[PIPELINE] ✅ Stage 2 完成: %d fragments, %d total rules", len(fragments), total_rules)
        yield {"stage": 2, "status": "done",
               "message": f"Stage 2 完成: {len(fragments)} 个单元 → {total_rules} 条规则",
               "elapsed": round(t2 - t1b, 1), "stage_total": 3,
               "detail": f"{len(fragments)} 个单元 → {total_rules} 条规则",
               "total_rules": total_rules}

        # --- Stage 3 ---
        yield {"stage": 3, "status": "running", "message": "Stage 3: 合并去重 + 审计报告...",
               "elapsed": 0, "stage_total": 3, "detail": "去重 + 分配 ID + 生成审计报告"}
        result = stage3_merge_and_audit(fragments, semantic_index, parsed_doc)
        t3 = time.time()
        n_rules = len(result["ir"]["rules"])

        ir_id = str(uuid.uuid4())[:8]
        output = {
            "ir_id": ir_id, "prd_id": prd_id,
            "yaml_content": json.dumps(result["ir"], ensure_ascii=False, indent=2),
            "ir_json": result["ir"],
            "audit": result["audit"],
            "audit_report": result["audit_report"],
            "skill_used": "default",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "pipeline_stats": {
                "stage1_seconds": round(t1 - t0, 1),
                "stage2_seconds": round(t2 - t1b, 1),
                "stage3_seconds": round(t3 - t2, 1),
                "total_seconds": round(t3 - t_total_start, 1),
            },
        }
        self._store[ir_id] = output

        # Save to disk
        output_dir = settings.OUTPUT_DIR / prd_id
        output_dir.mkdir(parents=True, exist_ok=True)
        with open(output_dir / f"ir_{ir_id}.json", "w", encoding="utf-8") as f:
            json.dump(output, f, ensure_ascii=False, indent=2, default=str)

        yield {"stage": 3, "status": "done", "message": f"流水线完成！共 {n_rules} 条规则",
               "elapsed": round(t3 - t2, 1), "stage_total": 3,
               "detail": f"合并后 {n_rules} 条规则，审计问题 {result['audit'].get('issue_count', 0)} 个",
               "done": True, "result": output}

    def _extract_one_sync(self, template: str, pkg: dict, client) -> list[dict]:
        """Synchronous rule extraction for one unit (runs in thread pool)."""
        prompt = _build_stage2_prompt(template, pkg)
        for attempt in range(3):
            try:
                raw = client.chat(
                    model=router.text_model,
                    messages=[
                        {"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON 数组，不输出任何其他文字。"},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.1,
                )
                json_str = extract_json_from_response(raw, expect_array=True)
                rules = json.loads(json_str)
                if isinstance(rules, list):
                    return rules
            except Exception as e:
                if attempt == 2:
                    raise
                time.sleep(2)
        return []

    async def run(self, prd_id: str, parsed_doc: dict) -> dict:
        """Run the full 3-stage pipeline and return the final IR + audit."""
        logger.info("[PIPELINE] ═══ 流水线启动 for PRD %s", prd_id)

        # Stage 1
        t0 = time.time()
        semantic_index = await stage1_semantic_index(parsed_doc)
        t1 = time.time()
        logger.info("[PIPELINE] Stage 1 耗时: %.1fs", t1 - t0)

        # Stage 2
        fragments = await stage2_extract_rules(semantic_index, parsed_doc)
        t2 = time.time()
        logger.info("[PIPELINE] Stage 2 耗时 %.1fs", t2 - t1)

        # Stage 3
        result = stage3_merge_and_audit(fragments, semantic_index, parsed_doc)
        t3 = time.time()
        logger.info("[PIPELINE] Stage 3 耗时 %.1fs", t3 - t2)

        ir_id = str(uuid.uuid4())[:8]
        output = {
            "ir_id": ir_id,
            "prd_id": prd_id,
            "yaml_content": json.dumps(result["ir"], ensure_ascii=False, indent=2),
            "ir_json": result["ir"],
            "audit": result["audit"],
            "audit_report": result["audit_report"],
            "skill_used": "default",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "pipeline_stats": {
                "stage1_seconds": round(t1 - t0, 1),
                "stage2_seconds": round(t2 - t1, 1),
                "stage3_seconds": round(t3 - t2, 1),
                "total_seconds": round(t3 - t0, 1),
            },
        }
        self._store[ir_id] = output

        # Save output
        output_dir = settings.OUTPUT_DIR / prd_id
        output_dir.mkdir(parents=True, exist_ok=True)
        with open(output_dir / f"ir_{ir_id}.json", "w", encoding="utf-8") as f:
            json.dump(output, f, ensure_ascii=False, indent=2, default=str)

        logger.info("[PIPELINE] ✅ 流水线完成: ir_id=%s, rules=%d, total=%.1fs",
                     ir_id, len(result["ir"]["rules"]), t3 - t0)
        return output

    async def get_ir(self, ir_id: str) -> dict | None:
        return self._store.get(ir_id)


ir_pipeline = IRPipeline()