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

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

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

import math
import time

from .. import config, evidence_model
from ..mcp_client import result_images, result_json

_LABEL_OVERALL = "overall"


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=0.02):
    """缩放到窗口、等待渲染、截一张图，返回 (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 capture_images(ctx, window, *, tiles=0, max_images=None):
    """采集视觉证据：整体一张 + 局部 tiles 张。

    - 整体 label="overall"，局部 label="tile-1"..；行优先。
    - 结束时**复位到整体视口**，保证后续流程从整幅视图起手。
    - 受 max_images（默认 config.EVIDENCE_MAX_IMAGES）截断。
    """
    if max_images is None:
        max_images = config.EVIDENCE_MAX_IMAGES
    images = []

    def _add(shot, label):
        if shot and len(images) < max_images:
            images.append(evidence_model.make_image(shot[0], shot[1], label))

    # 整体
    _add(_zoom_and_shot(ctx, window), _LABEL_OVERALL)

    # 局部（等分象限）
    for i, sub in enumerate(sub_windows(window, tiles), 1):
        if len(images) >= max_images:
            break
        _add(_zoom_and_shot(ctx, sub), f"tile-{i}")

    # 复位整体视口（不截图）
    ctx.tools.call("zoom_to_window", dict(window, margin=0.02))
    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. 视觉：整体 + 可选局部（多图）
    images = capture_images(ctx, win, tiles=config.EVIDENCE_DETAIL_TILES)

    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 源；PDF / 图片源预留（未来在对应分支接入离线解析 / 直接读图，
    产出与 CAD **同形**的 evidence：window=None、结构化键给空信封、images 带 label）。
    """
    source = evidence_model.source_type_of(unit)
    if source == evidence_model.SOURCE_CAD:
        return collect_frame(ctx, index, unit)
    # TODO(pdf/image): 接入离线解析（DWG/PDF）与直接读图。
    return None