cad-agent/dwg_review/__main__.py

"""CLI 入口：python -m dwg_review [options]"""

import argparse
import json
import sys
from pathlib import Path

from . import config
from .annotator import annotate, clear
from .checklist import build_system_prompt, checklist_summary
from .context import ReviewContext
from .mcp_client import McpError, McpToolError, result_json
from .phases import evidence as evidence_phase
from .phases import planning, report, review


def _dump_evidence(evidences):
    out_dir = Path(config.REPORT_DIR)
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / "evidence.json"
    trimmed = [{k: v for k, v in e.items() if k != "image_b64"} for e in evidences]
    path.write_text(json.dumps(trimmed, ensure_ascii=False, indent=2), encoding="utf-8")
    return path


def cmd_list_tools(ctx):
    tools = ctx.tools.refresh()
    print(f"共 {len(tools)} 个工具：")
    for tool in tools:
        print(f"  - {tool['name']}")
    return 0


def cmd_check(ctx):
    if not ctx.client.is_available():
        print(f"[X] {config.HOST}:{config.PORT} 未监听，请先启动已加载插件的 CAD。")
        return 1
    print(f"[OK] {config.HOST}:{config.PORT} 在监听")
    info = result_json(ctx.tools.call("get_drawing_info", {}), {})
    print("图纸信息:", json.dumps(info, ensure_ascii=False)[:400])
    return 0


def run_review(ctx, args):
    print("=== 图纸审查 Agent ===")
    print(f"清单：{checklist_summary(ctx.checklist)}")
    if ctx.llm is not None:
        print(f"模型：{ctx.llm.model}")
    print()

    print("[1/4] 规划：识别图框与概览 ...")
    plan = planning.plan(ctx)
    frames = plan.get("frames", [])
    print(f"      识别到 {len(frames)} 个图框")
    if not frames:
        print("未识别到任何图框，退出。")
        return 1

    system_prompt = build_system_prompt(ctx.checklist, allow_tools=True)

    print("[2/4] 取证：逐图框采集（串行） ...")
    evidences = []
    for idx, frame in enumerate(frames, 1):
        print(
            f"  - 图框 {idx}/{len(frames)} kind={frame.get('kind')} "
            f"sheet={frame.get('sheet_size')}"
        )
        ev = evidence_phase.collect_frame(ctx, idx, frame)
        if ev is None:
            print("    [跳过] 缺少有效包围盒")
            continue
        evidences.append(ev)
        print(
            f"    图元 {(ev.get('entities') or {}).get('total')} 条，"
            f"截图 {'有' if ev.get('image_b64') else '无'}"
        )

    if args.plan_only:
        path = _dump_evidence(evidences)
        print(f"[evidence] 已写入 {path}")
        return 0

    print("[3/4] 审查：逐图框 LLM 审查（含受控只读取证子循环） ...")
    findings = []
    for ev in evidences:
        print(f"  - 审查图框 {ev['index']} ...")
        text = review.review_frame(ctx, ev, system_prompt)
        findings.append(
            {
                "frame": ev["index"],
                "sheet": (ev.get("frame") or {}).get("sheet_size"),
                "finding": text,
            }
        )
        preview = text if len(text) <= 200 else text[:200] + "..."
        print(f"    结论：{preview}")

    print("[4/4] 汇总报告 ...")
    markdown = report.summarize(ctx, findings)
    issues = report.extract_issues(ctx, findings)
    report.save_report(markdown, findings, plan, issues)
    print()
    print(markdown)

    if args.annotate:
        print()
        print(f"[批注回写] {'(dry-run) ' if args.dry_run else ''}结构化问题 {len(issues)} 条")
        created, highlighted = annotate(ctx, issues, dry_run=args.dry_run)
        if not args.dry_run and (created or highlighted):
            hint = "python -m dwg_review --clear"
            if created:
                hint += f" {','.join(created)}"
            if highlighted:
                hint += f" --restore {','.join(highlighted)}"
            print(f"  撤销：{hint}")
    return 0


def build_parser():
    parser = argparse.ArgumentParser(
        prog="dwg_review", description="CAD 图纸审查 Agent（基于 cad_mcp_frame MCP 基座）"
    )
    parser.add_argument("--list-tools", action="store_true", help="列出基座注册的全部工具")
    parser.add_argument("--check", action="store_true", help="检查基座连通性并打印图纸概览")
    parser.add_argument("--plan-only", action="store_true", help="只做规划+取证，落盘 evidence.json")
    parser.add_argument("--frames", type=int, default=None, help="最多审查的图框数（0=全部）")
    parser.add_argument("--annotate", action="store_true", help="审查后把问题批注回写到图上")
    parser.add_argument("--dry-run", action="store_true", help="批注只打印不回写")
    parser.add_argument("--clear", default=None, help="删除给定 handle（逗号分隔）以撤销批注")
    parser.add_argument("--restore", default=None, help="还原给定 handle（逗号分隔）的原始颜色")
    parser.add_argument("--checklist", default=None, help="审查清单 JSON 路径")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)

    if args.checklist:
        config.CHECKLIST_PATH = Path(args.checklist)

    need_llm = not (
        args.list_tools or args.check or args.plan_only or args.clear or args.restore
    )
    try:
        ctx = ReviewContext.create(need_llm=need_llm)
    except McpError as exc:
        print(f"[X] 无法连接 CAD 基座：{exc}")
        return 1
    except Exception as exc:  # noqa: BLE001
        print(f"[X] 初始化失败：{exc}")
        return 1

    if args.frames is not None:
        ctx.max_frames = args.frames

    try:
        if args.list_tools:
            return cmd_list_tools(ctx)
        if args.check:
            return cmd_check(ctx)
        if args.clear or args.restore:
            handles = [h.strip() for h in (args.clear or "").split(",") if h.strip()]
            highlighted = [h.strip() for h in (args.restore or "").split(",") if h.strip()]
            clear(ctx, handles, highlighted)
            return 0
        return run_review(ctx, args)
    except (McpError, McpToolError) as exc:
        print(f"[X] 调用失败：{exc}")
        return 1
    except KeyboardInterrupt:
        print("\n[中断]")
        return 130


if __name__ == "__main__":
    sys.exit(main())