from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any


try:
    import fitz  # type: ignore
except Exception as exc:  # pragma: no cover
    raise SystemExit("PyMuPDF (fitz) is required: pip install pymupdf") from exc


_ELEM_RE = re.compile(r"(?m)^\s*([A-L]|[DEFGHIJK]\d{1,2})\b")
_ELEM_STRICT_RE = re.compile(r"^\s*([A-L]|[DEFGHIJK]\d{1,2})\b(?:\s*[\.\-–—:]\s*|\s+)")
_E_SUBSECTION_STRICT_RE = re.compile(r"^\s*(E[1-9])\b(?:\s*[\.\-–—:]\s*|\s+)")
_CONTENTS_WORD_RE = re.compile(r"(?i)\bcontents\b")
_PHOTO_REF_RE = re.compile(
    r"(?i)\b(photo(?:graph)?|figure|fig\.?)\s*[-:]?\s*(\d{1,3}|[A-Z])\b"
)


@dataclass(frozen=True, slots=True)
class ImageHit:
    page: int
    y0: float
    y1: float
    w: float
    h: float


def _iter_pdfs(folder: Path) -> list[Path]:
    return sorted([p for p in folder.rglob("*.pdf") if p.is_file()])


def _page_images(page: Any) -> list[ImageHit]:
    """Detect image blocks by parsing page dict blocks (type=1) + lightweight filters."""
    hits: list[ImageHit] = []
    ph = float(page.rect.height)
    pw = float(page.rect.width)
    try:
        blocks = page.get_text("dict").get("blocks", [])
    except Exception:
        blocks = []

    for b in blocks:
        if b.get("type") != 1:
            continue
        bbox = b.get("bbox") or [0, 0, 0, 0]
        x0, y0, x1, y1 = (float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3]))
        w = max(0.0, x1 - x0)
        h = max(0.0, y1 - y0)
        area = w * h

        # Filter tiny decorations and likely header/footer logos
        if area < 1500 or w < 30 or h < 30:
            continue
        if y0 < 60 or y1 > (ph - 60):
            continue
        # Filter narrow sidebars that are unlikely to be photos
        if pw > 0 and (w / pw) < 0.12 and h < 220:
            continue
        hits.append(ImageHit(page=int(page.number), y0=y0, y1=y1, w=w, h=h))
    return hits


def _is_likely_contents_page(text: str, headings_found: int) -> bool:
    """Heuristic: TOC pages list many codes and often include the word 'Contents'."""
    if headings_found >= 18:
        return True
    if _CONTENTS_WORD_RE.search(text) and headings_found >= 8:
        return True
    return False


def _page_headings(page: Any) -> list[tuple[float, str]]:
    """Return list of (y, code) for heading-like lines beginning with element code.

    This tries to avoid false positives from TOC tables by requiring:
    - the code token is followed by whitespace (e.g. 'E2 Roof coverings', not a table cell),
    - and the line uses a slightly larger font (typical of headings).
    """
    headings: list[tuple[float, str]] = []
    page_text_for_toc = ""
    found = 0
    try:
        blocks = page.get_text("dict").get("blocks", [])
    except Exception:
        blocks = []
    for b in blocks:
        if b.get("type") != 0:
            continue
        for line in b.get("lines", []) or []:
            spans = line.get("spans", []) or []
            if not spans:
                continue
            text = " ".join((s.get("text") or "") for s in spans).strip()
            if not text:
                continue
            page_text_for_toc += text + "\n"

            # Strict section/element heading match
            upper = text.upper()
            m = _ELEM_STRICT_RE.match(upper)
            if not m:
                continue
            code = m.group(1).upper()

            # For E-subsections, require E1..E9 (not just 'E')
            if code == "E":
                m2 = _E_SUBSECTION_STRICT_RE.match(upper)
                if not m2:
                    continue
                code = m2.group(1).upper()

            # Require a "heading-ish" font size to avoid TOC cells
            try:
                max_size = max(float(s.get("size") or 0.0) for s in spans)
            except Exception:
                max_size = 0.0
            if max_size and max_size < 10.5:
                continue

            y = float(line.get("bbox", [0, 0, 0, 0])[1])
            headings.append((y, code))
            found += 1

    if _is_likely_contents_page(page_text_for_toc, found):
        return []
    headings.sort(key=lambda t: t[0])
    return headings


def scan_pdf(fp: Path) -> dict[str, Any]:
    doc = fitz.open(str(fp))
    page_count = int(doc.page_count)
    per_code_seen: dict[str, int] = {}
    per_code_has_images: dict[str, int] = {}
    photo_refs: dict[str, int] = {}

    for pno in range(page_count):
        page = doc[pno]
        headings = _page_headings(page)
        images = _page_images(page)

        # Heading coverage (count codes that appear on a page at least once)
        codes_on_page = {c for _y, c in headings}
        for c in codes_on_page:
            per_code_seen[c] = per_code_seen.get(c, 0) + 1

        # Assign each image to nearest preceding heading on the page
        if headings and images:
            for img in images:
                best: str | None = None
                for hy, code in headings:
                    if hy <= img.y0 + 5:
                        best = code
                    else:
                        break
                if best:
                    per_code_has_images[best] = per_code_has_images.get(best, 0) + 1

        # Photo reference patterns from plain text (best-effort)
        try:
            txt = page.get_text() or ""
        except Exception:
            txt = ""
        for m in _PHOTO_REF_RE.finditer(txt):
            key = f"{m.group(1).lower()} {m.group(2)}"
            photo_refs[key] = photo_refs.get(key, 0) + 1

    # Normalise to ratios per code using page-level seen counts as denominator
    out_codes: dict[str, Any] = {}
    for code, seen_pages in per_code_seen.items():
        out_codes[code] = {
            "seen_pages": seen_pages,
            "images_assigned": int(per_code_has_images.get(code, 0)),
        }
    top_photo_refs = sorted(photo_refs.items(), key=lambda kv: kv[1], reverse=True)[:30]

    doc.close()

    return {
        "file": str(fp),
        "pages": page_count,
        "codes": out_codes,
        "top_photo_refs": top_photo_refs,
    }


def aggregate(scans: list[dict[str, Any]]) -> dict[str, Any]:
    # For each code: how many files show it, and how many files assign at least one image to it.
    files_seen: dict[str, int] = {}
    files_with_images: dict[str, int] = {}
    for s in scans:
        codes = s.get("codes") or {}
        for code, row in codes.items():
            files_seen[code] = files_seen.get(code, 0) + 1
            if int(row.get("images_assigned") or 0) > 0:
                files_with_images[code] = files_with_images.get(code, 0) + 1

    ratios = []
    for code, seen in files_seen.items():
        has = files_with_images.get(code, 0)
        ratios.append((has / seen if seen else 0.0, has, seen, code))
    ratios.sort(reverse=True)

    # Pull out E1-E9 specifically
    e_codes = [f"E{i}" for i in range(1, 10)]
    e_summary = []
    for code in e_codes:
        seen = files_seen.get(code, 0)
        has = files_with_images.get(code, 0)
        e_summary.append(
            {"code": code, "files_seen": seen, "files_with_images": has, "ratio": (has / seen if seen else 0.0)}
        )

    return {
        "files": len(scans),
        "by_code_sorted": [{"code": code, "files_with_images": has, "files_seen": seen, "ratio": r} for r, has, seen, code in ratios],
        "E1_E9": e_summary,
    }


def main() -> None:
    root = Path(__file__).resolve().parents[1]
    folders = [
        root / "Behrang RICS Documents",
        root / "RAW Context",
    ]
    payload: dict[str, Any] = {"root": str(root), "folders": {}}

    for folder in folders:
        pdfs = _iter_pdfs(folder)
        scans = [scan_pdf(fp) for fp in pdfs]
        payload["folders"][folder.name] = {
            "pdfs": [str(p) for p in pdfs],
            "aggregate": aggregate(scans),
            "files": scans,
        }

    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()