document_analyzer/skills/ir_generation_skill/main.py

"""
IR Generation Pipeline Orchestrator.

Run all four stages sequentially:
  python main.py [--skip-step1] [--skip-step2] [--skip-step2.5] [--skip-step3] [--test-only]

The pipeline reads the parsed PRD JSON from doc_parser and produces:
  - ir_final.json:      the final IR rules
  - ir_audit_report.md: completeness audit report for human review
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

import config

BASE_DIR = Path(__file__).parent


def _subprocess_env(extra: dict | None = None) -> dict:
    """Build environment dict for subprocesses, carrying forward overrides."""
    env = os.environ.copy()
    env.update(extra or {})
    return env


def run_step(script_name: str, description: str, extra_env: dict | None = None) -> bool:
    """Run a single pipeline step script, return True if it succeeded."""
    print(f"\n{'#' * 60}")
    print(f"# {description}")
    print(f"{'#' * 60}")
    script_path = BASE_DIR / script_name
    if not script_path.exists():
        print(f"错误: 脚本不存在 {script_path}")
        return False
    result = subprocess.run(
        [sys.executable, str(script_path)],
        cwd=str(BASE_DIR),
        env=_subprocess_env(extra_env),
    )
    return result.returncode == 0


def run_test(test_name: str, description: str, extra_env: dict | None = None) -> bool:
    """Run a test script, return True if all tests passed."""
    print(f"\n{'='*60}")
    print(f"测试: {description}")
    print(f"{'='*60}")
    test_path = BASE_DIR / "tests" / test_name
    if not test_path.exists():
        print(f"错误: 测试脚本不存在 {test_path}")
        return False
    result = subprocess.run(
        [sys.executable, str(test_path)],
        cwd=str(BASE_DIR),
        env=_subprocess_env(extra_env),
    )
    return result.returncode == 0


def main():
    parser = argparse.ArgumentParser(description="IR Generation Pipeline")
    parser.add_argument("--skip-step1", action="store_true",
                        help="跳过阶段一（语义索引）")
    parser.add_argument("--skip-step2", action="store_true",
                        help="跳过阶段二（IR 提取）")
    parser.add_argument("--skip-step2.5", "--skip-step2-5", action="store_true",
                        dest="skip_step2_5",
                        help="跳过阶段2.5（分支覆盖自动补全）")
    parser.add_argument("--skip-step3", action="store_true",
                        help="跳过阶段三（合并与审计）")
    parser.add_argument("--test-only", action="store_true",
                        help="仅运行测试，不调用 LLM")
    parser.add_argument(
        "--input", "-i", type=str, default=None,
        help="输入 JSON 文件路径（覆盖默认的 doc_parser 输出）"
    )
    parser.add_argument(
        "--provider", "-p", type=str, default=None,
        help="LLM provider: deepseek | dashscope（覆盖 IR_PROVIDER 环境变量）"
    )
    args = parser.parse_args()

    # Build extra env vars for subprocesses
    extra_env = {}
    if args.input:
        extra_env["IR_INPUT_JSON"] = args.input
        print(f"输入文件: {args.input}")
    if args.provider:
        extra_env["IR_PROVIDER"] = args.provider
        print(f"LLM Provider: {args.provider}")

    if args.test_only:
        all_ok = True
        all_ok &= run_test("test_step1.py", "Step 1 验证", extra_env)
        all_ok &= run_test("test_step2.py", "Step 2 验证", extra_env)
        all_ok &= run_test("test_step2_5.py", "Step 2.5 验证", extra_env)
        all_ok &= run_test("test_step3.py", "Step 3 验证", extra_env)
        sys.exit(0 if all_ok else 1)

    failures = []

    # Stage 1
    if not args.skip_step1:
        ok = run_step("step1_semantic_index.py",
                       "阶段一：宏观语义索引", extra_env)
        if not ok:
            failures.append("阶段一")
            print("\n阶段一失败，停止流水线。修复后重试。")
            sys.exit(1)
        run_test("test_step1.py", "Step 1 验证", extra_env)

    # Stage 2
    if not args.skip_step2:
        ok = run_step("step2_ir_extraction.py",
                       "阶段二：逐功能单元 IR 提取", extra_env)
        if not ok:
            failures.append("阶段二")
            print("\n阶段二失败，停止流水线。修复后重试。")
            sys.exit(1)
        run_test("test_step2.py", "Step 2 验证", extra_env)

    # Stage 2.5
    if not args.skip_step2_5:
        ok = run_step("step2_5_branch_coverage.py",
                       "阶段2.5：分支覆盖自动补全", extra_env)
        if not ok:
            failures.append("阶段2.5")
            print("\n阶段2.5失败，停止流水线。修复后重试。")
            sys.exit(1)
        run_test("test_step2_5.py", "Step 2.5 验证", extra_env)

    # Stage 3
    if not args.skip_step3:
        ok = run_step("step3_merge_and_audit.py",
                       "阶段三：确定性合并与完整性校验", extra_env)
        if not ok:
            failures.append("阶段三")
            sys.exit(1)
        run_test("test_step3.py", "Step 3 验证", extra_env)

    if failures:
        print(f"\n失败阶段: {', '.join(failures)}")
        sys.exit(1)

    print(f"\n{'='*60}")
    print("流水线全部完成!")
    print(f"最终 IR: {config.IR_FINAL_JSON}")
    print(f"审计报告: {config.IR_AUDIT_REPORT_MD}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()