"""非 CAD 源适配：PDF / 图片图纸的**离线**取证（不连 CAD）。

设计要点
--------
1. **世界坐标系**（两种源统一，且都 y 向上，与 `heatmap` / `tiling` 约定一致）：
   - PDF ：页内**点**坐标，原点左下，已应用 `/Rotate`（即 `page.get_size()` 的朝向）；
   - 图片：**像素**坐标，原点左下。
   故 `window` / `bbox` 与 CAD 侧同形，可直接喂给 `tiling` / `heatmap` 的纯函数。

2. **瓦片规划复用 CAD 侧链路**：把整页渲染（图片则读入）成一张缩略图 → 灰度二值化求
   墨迹密度 → 量化成与 `heatmap.build_grid` 同形的占用网格 → `heatmap.cluster_cells`
   聚类 → `tiling.plan_regions` 出瓦片窗口。**无需实体数据**，只靠像素墨迹。

3. **每张瓦片的渲染倍率由瓦片世界尺寸反推**：`zoom = IMAGE_MAX_EDGE / max(w,h)`，
   再夹到源允许的上限（PDF `PDF_TILE_MAX_ZOOM`；图片 `IMAGE_MAX_ZOOM`=1，**不放大**）。
   这与 `plan_regions` 的 `s_text` / `s_line` 是同一套数学：规划出的文字簇瓦片恰好
   落在 `zoom_text = TEXT_READABLE_PX / (TEXT_HEIGHT_MM · PT_PER_MM)`。

4. **文本层是次要的**：国标图纸的文字多为 SHX 描边（导出后变成线条），文本层几乎为空
   （实测某暖通图纸：正文页 0 个文本矩形）。仅当某页确有文本层（目录页 / 表格）时
   best-effort 抽取，塞进 `texts` 信封；主力识别交给 VLM 看渲染图。

5. `pypdfium2` **惰性 import**：只在真的处理 PDF 时导入，CAD / 图片路径零额外依赖。
"""

import base64
import io
import math
import os

from . import config, evidence_model, heatmap, tiling

_LABEL_OVERALL = "overall"

_PDF_EXTS = frozenset({".pdf"})
_IMAGE_EXTS = frozenset(
    {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp", ".gif"}
)
_CAD_EXTS = frozenset({".dwg", ".dxf", ".dwt"})


# ---------------------------------------------------------------------------
#  来源判定
# ---------------------------------------------------------------------------
def detect_source(path):
    """按扩展名判定来源类型（`cad` / `pdf` / `image`）；未知格式抛 ValueError。"""
    ext = os.path.splitext(str(path))[1].lower()
    if ext in _PDF_EXTS:
        return evidence_model.SOURCE_PDF
    if ext in _IMAGE_EXTS:
        return evidence_model.SOURCE_IMAGE
    if ext in _CAD_EXTS:
        return evidence_model.SOURCE_CAD
    raise ValueError(
        f"无法识别的图纸格式：{path}（支持 .dwg/.dxf/.pdf 与常见图片格式）"
    )


def is_file_source(source):
    """是否为文件源（非 CAD）。"""
    return source in (evidence_model.SOURCE_PDF, evidence_model.SOURCE_IMAGE)


# ---------------------------------------------------------------------------
#  PDF 坐标换算（纯函数，可单测）
# ---------------------------------------------------------------------------
def rotate_point(x, y, media_w, media_h, rotation):
    """PDF **未旋转**坐标 → **已旋转**（显示）坐标。

    `media_w`/`media_h` 为 MediaBox 宽高（未旋转）。rotation 取 0/90/180/270。
    注：仅 `/Rotate 90` 经实测验证（墨迹命中率 0.335 vs 其它假设 ≤0.064，并目视确认）。
    """
    r = int(rotation) % 360
    if r == 90:
        return (y, media_w - x)
    if r == 180:
        return (media_w - x, media_h - y)
    if r == 270:
        return (media_h - y, x)
    return (x, y)


def rect_to_world(rect, media_w, media_h, rotation):
    """未旋转 rect(l,b,r,t) → 已旋转世界 bbox(minx,miny,maxx,maxy)。"""
    x0, y0, x1, y1 = (float(v) for v in rect)
    pts = [
        rotate_point(x0, y0, media_w, media_h, rotation),
        rotate_point(x0, y1, media_w, media_h, rotation),
        rotate_point(x1, y0, media_w, media_h, rotation),
        rotate_point(x1, y1, media_w, media_h, rotation),
    ]
    xs = [p[0] for p in pts]
    ys = [p[1] for p in pts]
    return (min(xs), min(ys), max(xs), max(ys))


def crop_of(box, page_w, page_h):
    """世界 bbox（y 向上）→ `page.render(crop=...)` 的 (left, bottom, right, top) 裁边量。

    `crop` 是**要裁掉的边距**（不是绝对矩形），单位为点，且在**旋转后**的页面上生效。
    """
    x0, y0, x1, y1 = (float(v) for v in box)
    return (
        max(0.0, x0),
        max(0.0, y0),
        max(0.0, float(page_w) - x1),
        max(0.0, float(page_h) - y1),
    )


def guess_sheet_size(w_pt, h_pt, tol=0.02):
    """由页尺寸（点）反查标准幅面名（A0~A4，容差 2%）；不匹配返回 None。"""
    w_mm, h_mm = float(w_pt) / config.PT_PER_MM, float(h_pt) / config.PT_PER_MM
    longest, shortest = max(w_mm, h_mm), min(w_mm, h_mm)
    for name, (sw, sh) in tiling.SHEET_SIZES_MM.items():
        if abs(longest - sw) <= sw * tol and abs(shortest - sh) <= sh * tol:
            return name
    return None


# ---------------------------------------------------------------------------
#  墨迹分析（纯 PIL，无 numpy）
# ---------------------------------------------------------------------------
def _empty_grid(box=None):
    return {
        "window": tiling.box_to_window(box) if box else None,
        "cols": 0, "rows": 0, "cell_w": 0.0, "cell_h": 0.0,
        "counts": [], "coverage": [],
    }


def ink_grid(image, box, *, div=None, threshold=None, max_edge=None, min_frac=None):
    """把一张图按**墨迹密度**量化成占用网格（与 `heatmap.build_grid` 同形）。

    返回 `(grid, boxes)`：
      - `grid` ：可直接交给 `heatmap.cluster_cells`；
      - `boxes`：占用格的**世界 bbox** 列表（供 `tiling.content_predicate` 剔除空白瓦片）。

    步骤：灰度 → 缩到长边 ≤ `max_edge` → 按 `threshold` 二值化（暗 = 墨迹）→
    用 `Image.BOX` 块均值降到网格尺寸（每个像素即该格的墨迹面积占比）。

    `counts` 取 `int(frac * 100)`（≥1），故**墨迹越密计数越大**，`cluster_cells`
    的降序排列即是「内容密度降序」，密集区优先出图。
    """
    from PIL import Image

    div = config.HEATMAP_CELL_DIV if div is None else int(div)
    threshold = config.INK_THRESHOLD if threshold is None else int(threshold)
    max_edge = config.INK_ANALYSIS_MAX_EDGE if max_edge is None else int(max_edge)
    min_frac = config.INK_MIN_CELL_FRAC if min_frac is None else float(min_frac)

    box = tiling.norm_box(box)
    if box is None or box[2] <= box[0] or box[3] <= box[1] or image is None:
        return _empty_grid(box), []

    win = tiling.box_to_window(box)
    w, h = box[2] - box[0], box[3] - box[1]
    cell = heatmap.div_to_cell(win, div)
    if cell <= 0:
        return _empty_grid(box), []
    cols = max(1, int(math.ceil(w / cell)))
    rows = max(1, int(math.ceil(h / cell)))

    gray = image.convert("L")
    gw, gh = gray.size
    if max_edge > 0 and max(gw, gh) > max_edge:
        k = max_edge / float(max(gw, gh))
        gray = gray.resize(
            (max(1, int(round(gw * k))), max(1, int(round(gh * k)))), Image.BILINEAR
        )
    # 二值化：暗于阈值 → 墨迹（255）
    bw = gray.point(lambda p: 255 if p < threshold else 0, mode="L")
    # 块均值降采样到网格尺寸：每个输出像素 = 该格的墨迹面积占比 × 255
    small = bw.resize((cols, rows), Image.BOX)
    px = small.load()

    counts = [[0] * cols for _ in range(rows)]
    coverage = [[0.0] * cols for _ in range(rows)]
    for r in range(rows):
        for c in range(cols):
            frac = px[c, r] / 255.0
            coverage[r][c] = frac
            if frac >= min_frac:
                counts[r][c] = max(1, int(round(frac * 100)))

    grid = {
        "window": win, "cols": cols, "rows": rows,
        "cell_w": w / cols, "cell_h": h / rows,
        "counts": counts, "coverage": coverage,
    }
    boxes = [heatmap.cell_box(grid, r, c) for r, c in heatmap.occupied_cells(grid)]
    return grid, boxes


# ---------------------------------------------------------------------------
#  渲染倍率
# ---------------------------------------------------------------------------
def _tile_zoom(box, max_zoom, max_edge=None):
    """由瓦片世界尺寸反推渲染倍率：长边恰好落到 `max_edge` 像素，且不超过 `max_zoom`。

    与 `tiling.plan_regions` 的 `s_*` 是同一套数学（`zoom = 1/s`）：规划器把文字簇切到
    `s_text · IMAGE_MAX_EDGE` 世界宽度，这里就恰好还原成 `zoom_text`。
    """
    max_edge = config.IMAGE_MAX_EDGE if max_edge is None else int(max_edge)
    box = tiling.norm_box(box)
    if box is None:
        return float(max_zoom)
    longest = max(box[2] - box[0], box[3] - box[1])
    if longest <= 0 or max_edge <= 0:
        return float(max_zoom)
    return min(float(max_zoom), max_edge / longest)


def file_scales(source):
    """文件源的 `(s_text, s_line)`——CAD 图形单位 / 像素。

    PDF 的世界单位是**点**，且 PDF 内容本身就是**纸面 1:1**（出图比例已烘焙进图形），
    故等价于 `frame_scale = 1`、纸面字高 = `TEXT_HEIGHT_MM · PT_PER_MM` 点。
    图片的世界单位是**像素**：1 世界单位 = 1 像素，且不允许放大，故 s 恒为 1。
    """
    if source == evidence_model.SOURCE_IMAGE:
        return 1.0, 1.0
    text_pt = config.TEXT_HEIGHT_MM * config.PT_PER_MM
    return (
        tiling.world_per_px(1.0, text_pt, config.TEXT_READABLE_PX),
        tiling.world_per_px(1.0, text_pt, config.LINE_READABLE_PX),
    )


def snap_box(box, bounds=None):
    """把 bbox 对齐到**整数世界坐标**（图片源的世界单位就是源像素）。

    渲染出来的像素只能是整数，若不对齐，注册的 `window` 与实际像素会差最多 1px，
    使「窗口 ↔ 像素」映射不精确（未来按像素反算世界坐标会累积误差）。
    `bounds` 给定时顺带夹回边界；退化返回 None。
    """
    b = tiling.norm_box(box)
    if b is None:
        return None
    x0, y0 = math.floor(b[0]), math.floor(b[1])
    x1, y1 = math.ceil(b[2]), math.ceil(b[3])
    if bounds is not None:
        bb = tiling.norm_box(bounds)
        if bb is not None:
            x0, y0 = max(x0, math.floor(bb[0])), max(y0, math.floor(bb[1]))
            x1, y1 = min(x1, math.ceil(bb[2])), min(y1, math.ceil(bb[3]))
    if x1 <= x0 or y1 <= y0:
        return None
    return (float(x0), float(y0), float(x1), float(y1))


def _png_b64(img):
    buf = io.BytesIO()
    if img.mode not in ("RGB", "L"):
        img = img.convert("RGB")
    img.save(buf, format="PNG", optimize=True)
    return base64.b64encode(buf.getvalue()).decode("ascii")


# ---------------------------------------------------------------------------
#  读取器：PDF / 图片
# ---------------------------------------------------------------------------
class PdfReader:
    """PDF 源：按页取证（内存渲染 + best-effort 文本层）。惰性 import pypdfium2。"""

    def __init__(self, path):
        import pypdfium2 as pdfium  # 惰性：仅 PDF 路径需要

        self.path = str(path)
        self.doc = pdfium.PdfDocument(self.path)
        self.page_count = len(self.doc)
        self._metas = {}

    # -- 生命周期 ----------------------------------------------------------
    def close(self):
        try:
            self.doc.close()
        except Exception:  # noqa: BLE001
            pass

    # -- 页元数据 ----------------------------------------------------------
    def _meta(self, index):
        """(旋转后宽, 旋转后高, MediaBox 宽, MediaBox 高, rotation)，均以点计。"""
        m = self._metas.get(index)
        if m is None:
            page = self.doc[index]
            media = page.get_mediabox()
            mw = float(media[2]) - float(media[0])
            mh = float(media[3]) - float(media[1])
            size = page.get_size()  # 已应用 /Rotate
            rw, rh = float(size[0]), float(size[1])
            m = (rw, rh, mw, mh, int(page.get_rotation()))
            self._metas[index] = m
        return m

    # -- 规划 --------------------------------------------------------------
    def frames(self):
        """每页一个「图框」（bbox = 页内点坐标，y 向上，已应用 /Rotate）。"""
        out = []
        for i in range(self.page_count):
            rw, rh, _mw, _mh, _rot = self._meta(i)
            out.append(
                {
                    "bbox": [0.0, 0.0, rw, rh],
                    "kind": "page",
                    "page_index": i,
                    "sheet_size": guess_sheet_size(rw, rh),
                    "sheet_scale": 1.0,
                    "layer": None,
                    "confidence": 1.0,
                    "source_type": evidence_model.SOURCE_PDF,
                }
            )
        return out

    # -- 渲染 / 文本 -------------------------------------------------------
    def render(self, box, zoom, page_index=0):
        """把世界 bbox 渲染成 PIL 图（`crop` 为点、在旋转后页面上生效）。

        退化 bbox（宽或高 ≤ 0）会让 `crop` 的边距之和等于页尺寸，pdfium 直接报错，
        故在此提前拒绝。
        """
        box = tiling.norm_box(box)
        if box is None or box[2] <= box[0] or box[3] <= box[1]:
            return None
        rw, rh, _mw, _mh, _rot = self._meta(page_index)
        page = self.doc[page_index]
        bitmap = page.render(
            scale=float(zoom), rotation=0, crop=crop_of(box, rw, rh)
        )
        img = bitmap.to_pil()
        return img.convert("RGB") if img.mode not in ("RGB", "L") else img

    def text_items(self, index, box=None, min_chars=None, limit=0):
        """best-effort 抽文本层：逐文本矩形取串 + 世界 bbox。

        国标图纸文字多为 SHX 描边（导出成线条），多数页**整页 0 个文本矩形**——
        返回空列表属**预期**，不是 bug。真正的识别靠 VLM 看渲染图。
        """
        min_chars = config.PDF_TEXT_MIN_CHARS if min_chars is None else int(min_chars)
        _rw, _rh, mw, mh, rot = self._meta(index)
        page = self.doc[index]
        try:
            tp = page.get_textpage()
        except Exception:  # noqa: BLE001 - 无文本层 / 解析失败都按「无文本」处理
            return []
        try:
            total = tp.count_rects()  # get_rect 的前置调用（pdfium 要求）
            items = []
            for k in range(total):
                if limit and len(items) >= limit:
                    break
                try:
                    rect = tp.get_rect(k)
                    text = (tp.get_text_bounded(*rect) or "").strip()
                except Exception:  # noqa: BLE001 - 单个矩形失败不影响整页
                    continue
                if len(text) < min_chars:
                    continue
                world = rect_to_world(rect, mw, mh, rot)
                if box is not None and not tiling.boxes_intersect(world, box):
                    continue
                items.append({"content": text, "bbox": list(world)})
            return items
        finally:
            try:
                tp.close()
            except Exception:  # noqa: BLE001
                pass


class ImageReader:
    """图片源：整图即一页，按需裁剪（**只缩不放**，放大只会插值出假细节）。"""

    def __init__(self, path):
        from PIL import Image

        self.path = str(path)
        self.img = Image.open(self.path)
        self.img.load()
        self.width, self.height = self.img.size
        self.page_count = 1

    def close(self):
        try:
            self.img.close()
        except Exception:  # noqa: BLE001
            pass

    def frames(self):
        return [
            {
                "bbox": [0.0, 0.0, float(self.width), float(self.height)],
                "kind": "image",
                "page_index": 0,
                "sheet_size": None,
                "sheet_scale": 1.0,
                "layer": None,
                "confidence": 1.0,
                "source_type": evidence_model.SOURCE_IMAGE,
            }
        ]

    def render(self, box, zoom, page_index=0):
        from PIL import Image

        # 世界单位 = 源像素 → 对齐到整数，保证「窗口 ↔ 像素」精确对应
        box = snap_box(box, (0.0, 0.0, float(self.width), float(self.height)))
        if box is None:
            return None
        # 世界 y 向上 → 图像 y 向下
        left = int(box[0])
        upper = int(self.height - box[3])
        right = int(box[2])
        lower = int(self.height - box[1])
        if right <= left or lower <= upper:
            return None
        out = self.img.crop((left, upper, right, lower))
        if zoom and 0 < zoom < 1.0:
            out = out.resize(
                (
                    max(1, int(round(out.size[0] * zoom))),
                    max(1, int(round(out.size[1] * zoom))),
                ),
                Image.LANCZOS,
            )
        return out.convert("RGB") if out.mode not in ("RGB", "L") else out

    def text_items(self, index=0, box=None, **kwargs):
        """图片没有文本层。"""
        return []


# ---------------------------------------------------------------------------
#  读取器缓存（挂在 ctx 上，避免每页重复打开 PDF）
# ---------------------------------------------------------------------------
def open_reader(ctx):
    """打开（并缓存到 `ctx.source_reader`）文件源读取器。"""
    source = ctx.source
    if not is_file_source(source):
        raise ValueError(f"来源 {source} 不是文件源，无需读取器")
    reader = getattr(ctx, "source_reader", None)
    if reader is not None:
        return reader
    path = getattr(ctx, "source_path", None)
    if not path:
        raise ValueError("文件源未指定路径（ctx.source_path 为空）")
    reader = PdfReader(path) if source == evidence_model.SOURCE_PDF else ImageReader(path)
    ctx.source_reader = reader
    return reader


def close_reader(ctx):
    """释放文件源读取器（幂等）。"""
    reader = getattr(ctx, "source_reader", None)
    if reader is not None:
        reader.close()
        ctx.source_reader = None


# ---------------------------------------------------------------------------
#  规划（与 phases.planning 的 CAD 分支同形）
# ---------------------------------------------------------------------------
def plan_file(ctx):
    """文件源的「规划」：产出与 CAD 同形的 plan（每页 / 整图一个图框，保持原始页序）。"""
    reader = open_reader(ctx)
    frames = reader.frames()
    if ctx.source == evidence_model.SOURCE_PDF and config.PDF_MAX_PAGES > 0:
        frames = frames[: config.PDF_MAX_PAGES]
    if ctx.max_frames > 0:
        frames = frames[: ctx.max_frames]
    info = {
        "source_type": ctx.source,
        "path": ctx.source_path,
        "page_count": getattr(reader, "page_count", len(frames)),
        "unit": "pt" if ctx.source == evidence_model.SOURCE_PDF else "px",
    }
    return {
        "drawing_info": info,
        "layers": evidence_model.empty_envelope(),
        "layouts": evidence_model.empty_envelope(),
        "frames": frames,
    }


# ---------------------------------------------------------------------------
#  取证
# ---------------------------------------------------------------------------
def _same_box(a, b, tol=1e-6):
    """两个 bbox 是否（近似）相同。"""
    return a is not None and b is not None and all(
        abs(float(x) - float(y)) <= tol for x, y in zip(a, b)
    )


def _capture_file(ctx, frame, box, reader, source, text_boxes):
    """采集视觉证据：整体一张 + 自适应墨迹瓦片若干。"""
    page_index = frame.get("page_index", 0)
    is_pdf = source == evidence_model.SOURCE_PDF
    max_zoom = config.PDF_TILE_MAX_ZOOM if is_pdf else config.IMAGE_MAX_ZOOM
    overview_zoom = config.PDF_OVERVIEW_MAX_ZOOM if is_pdf else config.IMAGE_MAX_ZOOM

    images = []

    def _add(img, label, win):
        if img is None or len(images) >= config.EVIDENCE_MAX_IMAGES:
            return False
        images.append(
            evidence_model.make_image(
                _png_b64(img), "image/png", label, window=list(win)
            )
        )
        return True

    # 整体图：整页 / 整图渲染（长边 ≈ IMAGE_MAX_EDGE，且不超过 overview 上限）
    overall_zoom = _tile_zoom(box, overview_zoom)
    _add(reader.render(box, overall_zoom, page_index), _LABEL_OVERALL, box)

    if not config.TILING_ENABLED:
        return images

    # 墨迹分析：低倍缩略图 → 占用网格 → 簇 → 瓦片
    thumb = reader.render(
        box, _tile_zoom(box, max_zoom, config.INK_ANALYSIS_MAX_EDGE), page_index
    )
    grid, ink_boxes = ink_grid(thumb, box)
    clusters = heatmap.cluster_cells(grid, dilate=config.HEATMAP_CLUSTER_DILATE)
    if not clusters:
        return images

    s_text, s_line = file_scales(source)
    if s_text <= 0 or s_line <= 0:
        return images
    flags = tiling.cluster_text_flags(clusters, text_boxes)
    text_idx = {i for i, is_text in enumerate(flags) if is_text}

    regions = tiling.plan_regions(
        frame, clusters,
        text_clusters=text_idx, s_text=s_text, s_line=s_line,
        view_px=(config.IMAGE_MAX_EDGE, config.IMAGE_MAX_EDGE),
        max_tiles=config.TILE_MAX_COUNT, overlap=config.TILE_OVERLAP,
        keep=tiling.content_predicate(ink_boxes),
    )
    n = 0
    for win in regions:
        if len(images) >= config.EVIDENCE_MAX_IMAGES:
            break
        tile_box = tiling.window_to_box(win)
        if tile_box is None:
            continue
        if not is_pdf:
            # 图片源世界单位是源像素：先对齐整数，注册的 window 才与像素精确对应
            tile_box = snap_box(tile_box, box)
            if tile_box is None:
                continue
        zoom = _tile_zoom(tile_box, max_zoom)
        # 整页即一个簇时，瓦片窗口会与整体图重合且分辨率不更高 → 省掉这张重复图
        if _same_box(tile_box, box) and zoom <= overall_zoom + 1e-9:
            continue
        if _add(reader.render(tile_box, zoom, page_index), f"tile-{n + 1}", tile_box):
            n += 1
    return images


def collect(ctx, index, unit):
    """文件源取证入口；产出与 `phases.evidence.collect_frame` **同形**的 evidence。"""
    reader = open_reader(ctx)
    box = tiling.frame_bbox(unit)
    if box is None:
        return None
    source = evidence_model.source_type_of(unit)

    evidence = {
        "source_type": source,
        "index": index,
        "frame": unit,
        "window": tiling.box_to_window(box),
        "entities": evidence_model.empty_envelope(),
        "texts": evidence_model.empty_envelope(),
        "dimensions": evidence_model.empty_envelope(),
        "blocks": evidence_model.empty_envelope(),
        "polylines": evidence_model.empty_envelope(),
        "images": [],
    }

    # best-effort 文本层（仅 PDF 可能有；图片恒为空）
    page_index = unit.get("page_index", 0)
    items = reader.text_items(page_index, box, limit=ctx.frame_limit)
    evidence["texts"] = {"total": len(items), "items": items}
    text_boxes = [i["bbox"] for i in items if i.get("bbox")]

    evidence["images"] = _capture_file(ctx, unit, box, reader, source, text_boxes)
    return evidence