cad-agent/dwg_review/phases/evidence.py

"""阶段 2：逐审查单元取证（多源证据层 + 视觉多图）。

对每个审查单元：取结构化数据（实体/文字/标注/块/多段线）+ 视觉截图（整体 + 可选局部）。
结构化数据是权威依据，截图仅作视觉上下文。

多源设计：`collect()` 按审查单元来源分发；当前仅实现 CAD 源，PDF / 图片源预留 seam
（未来在对应分支产出**同形** evidence，下游无需改动）。
"""

import contextlib
import math
import time

from .. import config, evidence_model, heatmap, sources, tiling
from ..llm_client import image_size_b64
from ..mcp_client import result_images, result_json

_LABEL_OVERALL = "overall"

# 绘图区最大化用的工具名（旧版插件没有此工具 → 静默降级）
_TOOL_DRAWING_AREA = "set_drawing_area"

# zoom 时四周留边比例（与 `zoom_to_window` 的默认 margin 一致）
_ZOOM_MARGIN = 0.02


@contextlib.contextmanager
def maximized_drawing_area(ctx):
    """截屏前最大化绘图区（窗口铺满屏幕 + 干净屏幕），退出时还原。

    实测绘图区像素 2110x462 → 3310x1084（3.68x）；在同样的长边下采样预算下，
    有效像素更多、细节更清楚，自适应瓦片张数也更少。

    - 受 `config.DRAWING_AREA_MAXIMIZE` 控制；
    - 工具不存在（旧版插件）/ 调用失败 → **静默降级**：照常取证，只是绘图区小一些；
    - 退出时一定尝试还原；**还原失败不掩盖业务异常**（不 raise）。
    - 离线源（PDF / 图片）无 CAD 连接（`ctx.tools is None`）→ 直接放行。
    """
    if not config.DRAWING_AREA_MAXIMIZE or ctx.tools is None:
        yield
        return

    entered = False
    try:
        if ctx.tools.has(_TOOL_DRAWING_AREA):
            ctx.tools.call(_TOOL_DRAWING_AREA, {"mode": "max"})
            entered = True
            time.sleep(config.DRAWING_AREA_ENTER_WAIT)
    except Exception:  # noqa: BLE001 - 界面优化失败不应中断审查
        entered = False

    try:
        yield
    finally:
        if entered:
            try:
                ctx.tools.call(_TOOL_DRAWING_AREA, {"mode": "restore"})
                time.sleep(config.DRAWING_AREA_RESTORE_WAIT)
            except Exception:  # noqa: BLE001
                pass


def bbox_to_window(bbox):
    """[minx,miny,maxx,maxy] -> window 对象。"""
    return {
        "min_x": float(bbox[0]),
        "min_y": float(bbox[1]),
        "max_x": float(bbox[2]),
        "max_y": float(bbox[3]),
    }


def sub_windows(window, tiles):
    """把 window 等分为 tiles 个子窗口（tiles=4 → 2x2 象限）。

    tiles<=1 返回 []。返回 [{"min_x",..,"max_x",..}, ...]，行优先（左上→右上→左下→右下）。
    """
    if not tiles or tiles <= 1:
        return []
    cols = int(math.ceil(math.sqrt(tiles)))
    rows = int(math.ceil(tiles / cols))
    min_x, min_y = window["min_x"], window["min_y"]
    max_x, max_y = window["max_x"], window["max_y"]
    w = (max_x - min_x) / cols
    h = (max_y - min_y) / rows
    out = []
    for r in range(rows):
        for c in range(cols):
            if len(out) >= tiles:
                break
            out.append(
                {
                    "min_x": min_x + c * w,
                    "min_y": max_y - (r + 1) * h,
                    "max_x": min_x + (c + 1) * w,
                    "max_y": max_y - r * h,
                }
            )
    return out


def zoom_and_shot(ctx, window, margin=_ZOOM_MARGIN):
    """缩放到窗口、等待渲染、截一张图，返回 (b64, mime) 或 None。"""
    ctx.tools.call("zoom_to_window", dict(window, margin=margin))
    time.sleep(config.RENDER_WAIT)
    shot = ctx.tools.call("get_viewport_screenshot", {})
    images = result_images(shot)
    if not images:
        return None
    return images[0]


def _is_truncated(*envelopes):
    """任一信封 `total > len(items)` 即视为**被截断**（数据不全）。"""
    for env in envelopes:
        if not isinstance(env, dict):
            continue
        items = env.get("items")
        if not isinstance(items, list):
            continue
        try:
            total = int(env.get("total") or len(items))
        except (TypeError, ValueError):
            total = len(items)
        if total > len(items):
            return True
    return False


def _text_boxes(entities, texts, dimensions):
    """收集「文字类」实体 bbox。

    标准文字 / 标注信封（`get_text_content` / `get_dimensions`）全收；
    再加**类名含 text/label/dim 等关键词**的自定义实体（`tiling.is_text_type`）。
    """
    boxes = list(heatmap.entity_boxes(texts))
    boxes.extend(heatmap.entity_boxes(dimensions))
    items = entities.get("items") if isinstance(entities, dict) else None
    for it in items or []:
        if isinstance(it, dict) and tiling.is_text_type(it.get("type")):
            b = heatmap.norm_box(it.get("bbox"))
            if b is not None:
                boxes.append(b)
    return boxes


def _content_predicate(boxes):
    """构造「该窗口内是否有实体」谓词；无实体返回 None（委托 `tiling.content_predicate`）。"""
    return tiling.content_predicate(boxes)


def _apply_keep(regions, keep):
    """按 keep 过滤窗口；过滤后为空则**退回原集合**（宁可多拍也不漏）。"""
    if keep is None or not regions:
        return regions
    kept = [w for w in regions if keep(w)]
    return kept or regions


def _plan_tiles(window, frame, boxes, text_boxes, truncated, view_px):
    """尝试按实体分布规划瓦片窗口；任何条件不满足都返回 []（调用方退回固定策略）。

    降级条件：规划被关闭 / 实体数据被截断 / 缺图框 / 无实体 / 拿不到绘图区像素 /
    比例未知（`frame_scale` 返 0）。
    """
    if not config.TILING_ENABLED or truncated or not frame or not boxes or not view_px:
        return []
    scale = tiling.frame_scale(frame)
    s_text = tiling.world_per_px(scale, config.TEXT_HEIGHT_MM, config.TEXT_READABLE_PX)
    # 线条区只需看清走向：同以「最小关注特征尺寸」（纸面 TEXT_HEIGHT_MM）为基准，
    # 但像素门槛更低（R_line < R_text）→ 覆盖范围更大、瓦片更少。
    s_line = tiling.world_per_px(scale, config.TEXT_HEIGHT_MM, config.LINE_READABLE_PX)
    if s_text <= 0 or s_line <= 0:
        return []
    eff = tiling.effective_px(view_px, config.IMAGE_MAX_EDGE)
    if eff[0] <= 0 or eff[1] <= 0:
        return []

    cell = heatmap.div_to_cell(window, config.HEATMAP_CELL_DIV)
    grid = heatmap.build_grid(window, boxes, cell)
    clusters = heatmap.cluster_cells(grid, dilate=config.HEATMAP_CLUSTER_DILATE)
    if not clusters:
        return []
    flags = tiling.cluster_text_flags(clusters, text_boxes)
    text_idx = {i for i, is_text in enumerate(flags) if is_text}

    return tiling.plan_regions(
        frame, clusters,
        text_clusters=text_idx, s_text=s_text, s_line=s_line, view_px=eff,
        max_tiles=config.TILE_MAX_COUNT, overlap=config.TILE_OVERLAP,
        keep=_content_predicate(boxes),
    )


def capture_images(ctx, window, *, tiles=0, max_images=None, frame=None,
                   boxes=None, text_boxes=None, truncated=False):
    """采集视觉证据：整体一张 + 局部若干张（自适应瓦片优先，否则等分象限回退）。

    - 整体 label="overall"，局部 label="tile-1"..；行优先。
    - 局部窗口优先由 `_plan_tiles` 按**实体分布**给出（空白区不出图）；
      不可用时退回 `sub_windows(window, tiles)`，**同样按内容过滤**空白象限
      （过滤后为空则退回未过滤集合，保证不漏拍）。
    - 每张图登记其**实际可见世界范围**（`window` 字段）——审查阶段的 `inspect_region`
      靠它把「截图上的归一化坐标」映射回 WCS（逐级下钻）。
    - 结束时**复位到整体视口**，保证后续流程从整幅视图起手。
    - 受 max_images（默认 config.EVIDENCE_MAX_IMAGES）截断。
    """
    if max_images is None:
        max_images = config.EVIDENCE_MAX_IMAGES
    images = []

    # 整体图：既是证据，也用于测量绘图区**实际像素**（决定单张截图的世界覆盖范围）
    overall = zoom_and_shot(ctx, window)
    view_px = image_size_b64(overall[0]) if overall else None

    def _add(shot, label, win):
        """登记截图；`win` 为送给 CAD 的窗口，用于反算实际可见世界范围。"""
        if shot and len(images) < max_images:
            images.append(
                evidence_model.make_image(
                    shot[0], shot[1], label,
                    window=tiling.visible_box(
                        tiling.window_to_box(win), view_px, _ZOOM_MARGIN
                    ),
                )
            )

    _add(overall, _LABEL_OVERALL, window)

    # 内容过滤：空白区不出图。按**实体 bbox** 判定，规划与回退两条路径共用同一谓词；
    # 实体数据被截断时不可信，放弃过滤（宁可多拍也不漏）。
    keep = None if truncated else _content_predicate(boxes)

    # 局部：优先自适应瓦片，失败退回等分象限（回退路径同样按内容过滤）
    regions = _plan_tiles(window, frame, boxes, text_boxes, truncated, view_px)
    if not regions:
        regions = _apply_keep(sub_windows(window, tiles), keep)

    for i, sub in enumerate(regions, 1):
        if len(images) >= max_images:
            break
        _add(zoom_and_shot(ctx, sub), f"tile-{i}", sub)

    # 复位整体视口（不截图）
    ctx.tools.call("zoom_to_window", dict(window, margin=_ZOOM_MARGIN))
    return images


def collect_frame(ctx, index, frame):
    """采集单个图框（CAD 源）的证据；bbox 无效时返回 None。"""
    bbox = frame.get("bbox")
    if not bbox or len(bbox) != 4:
        return None

    win = bbox_to_window(bbox)
    limit = ctx.frame_limit

    # a. 结构化数据（权威；按 window 查询，不依赖缩放）
    entities = result_json(ctx.tools.call("query_entities", {"window": win, "limit": limit}), {})
    texts = result_json(ctx.tools.call("get_text_content", {"window": win, "limit": limit}), {})
    dimensions = result_json(ctx.tools.call("get_dimensions", {"window": win, "limit": limit}), {})
    blocks = result_json(
        ctx.tools.call("get_block_references", {"window": win, "limit": limit}), {}
    )
    polylines = result_json(
        ctx.tools.call("get_polylines_info", {"window": win, "limit": limit}), {}
    )

    # b. 视觉：整体 + 局部（优先按实体分布自适应瓦片；数据截断则退回固定策略）
    truncated = _is_truncated(entities, texts, dimensions, blocks, polylines)
    boxes = heatmap.drop_oversized(
        heatmap.entity_boxes(entities), win, config.HEATMAP_MAX_AREA_FRAC
    )
    images = capture_images(
        ctx, win,
        tiles=config.EVIDENCE_DETAIL_TILES,
        frame=frame,
        boxes=boxes,
        text_boxes=_text_boxes(entities, texts, dimensions),
        truncated=truncated,
    )

    return {
        "source_type": evidence_model.source_type_of(frame),
        "index": index,
        "frame": frame,
        "window": win,
        "entities": entities,
        "texts": texts,
        "dimensions": dimensions,
        "blocks": blocks,
        "polylines": polylines,
        "images": images,
    }


def collect(ctx, index, unit):
    """按审查单元来源分发取证。

    CAD 源走 `collect_frame`（CAD 工具 + 截屏）；PDF / 图片源走 `sources.collect`
    （内存渲染 + 墨迹瓦片），产出与 CAD **同形**的 evidence：`window` 有值、
    结构化键给空信封（PDF 另有 best-effort 文本层）、`images` 带 label。
    """
    source = evidence_model.source_type_of(unit)
    if source == evidence_model.SOURCE_CAD:
        return collect_frame(ctx, index, unit)
    if sources.is_file_source(source):
        return sources.collect(ctx, index, unit)
    return None