"""Figure extraction: image detection, caption association."""
from __future__ import annotations

import re
import fitz
from ..types import BBox, Figure, TextLine, DocumentBlock, BlockType
from ..utils.geometry import vertical_distance, horizontal_overlap


# ---------------------------------------------------------------------------
# Roman-numeral helpers
# ---------------------------------------------------------------------------

_ROMAN_PATTERN = re.compile(
    r"^M{0,4}(?:CM|CD|D?C{0,3})(?:XC|XL|L?X{0,3})(?:IX|IV|V?I{0,3})$",
    re.IGNORECASE,
)


def _roman_to_int(s: str) -> int | None:
    """Convert a Roman numeral string to int, or return None if invalid."""
    s = s.upper().strip()
    if not s or not _ROMAN_PATTERN.match(s):
        return None
    vals = {"I": 1, "V": 5, "X": 10, "L": 50, "C": 100, "D": 500, "M": 1000}
    result = 0
    prev = 0
    for ch in reversed(s):
        v = vals.get(ch, 0)
        if v < prev:
            result -= v
        else:
            result += v
        prev = v
    return result if result > 0 else None


# ---------------------------------------------------------------------------
# Caption detection regex
# ---------------------------------------------------------------------------

# Matches the START of a caption line.
#
# Two sub-patterns cover Arabic and Roman-numeral caption numbers:
#
#   _CAPTION_ARABIC  — "Figure 1: ..." / "Table 3.1. ..." / "Fig. 2 text"
#   _CAPTION_ROMAN   — "TABLE I. ..." / "Table IV: ..."
#
# Both return named groups: keyword, number, rest.
#
# Roman numerals require punctuation (. or :) or end-of-string after the
# number to avoid false positives on body sentences like "Figure I think…".

_CAPTION_ARABIC = re.compile(
    r"^(?P<keyword>Figure|Fig\.?|Table|Tab\.?)"  # keyword
    r"\s+"
    r"(?P<number>\d+(?:\.\d+)*)"                 # Arabic / compound number
    r"[.:\s]"                                     # separator (loose)
    r"(?P<rest>.*)",
    re.IGNORECASE,
)

_CAPTION_ROMAN = re.compile(
    r"^(?P<keyword>Figure|Fig\.?|Table|Tab\.?)"  # keyword
    r"\s+"
    r"(?P<number>[IVXLCDM]{1,10})"               # Roman numeral (upper/lower via flag)
    r"(?:[.:\s]|$)"                               # separator OR end-of-string (strict)
    r"(?P<rest>.*)",
    re.IGNORECASE,
)


def _match_caption(text: str):
    """Try to match text as a caption start.  Returns re.Match or None."""
    m = _CAPTION_ARABIC.match(text)
    if m:
        return m
    m = _CAPTION_ROMAN.match(text)
    if m:
        # Extra guard: validate that the number group is a real Roman numeral
        if _roman_to_int(m.group("number")) is not None:
            return m
    return None


def _parse_caption_number(num_str: str) -> int:
    """Return an integer caption number from Arabic or Roman numeral string.

    For compound numbers like '3.1' we use the first component.
    Returns 0 if unparseable.
    """
    num_str = num_str.strip()
    # Arabic (possibly compound)
    if re.match(r"^\d", num_str):
        return int(num_str.split(".")[0])
    # Roman numeral
    v = _roman_to_int(num_str)
    return v if v is not None else 0


def _caption_type(keyword: str) -> str:
    """Normalise caption keyword to 'figure' or 'table'."""
    kw = keyword.lower().rstrip(".")
    if kw in ("fig", "figure"):
        return "figure"
    return "table"


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def extract_vector_figures(
    page: fitz.Page,
    page_num: int,
    table_bboxes: list[BBox] | None = None,
) -> list[Figure]:
    """Detect vector graphics (architecture diagrams, flowcharts) on a page.

    Uses page.get_drawings() to find path-based drawings, clusters them
    spatially, and returns figures for clusters large enough to be diagrams.
    Table rules (horizontal/vertical lines matching table bboxes) are excluded.
    """
    try:
        drawings = page.get_drawings()
    except Exception:
        return []

    if not drawings:
        return []

    _table_bboxes = table_bboxes or []

    # Filter out drawing commands that overlap table regions
    def _in_table(rect: tuple[float, float, float, float]) -> bool:
        r = BBox(rect[0], rect[1], rect[2], rect[3])
        for tb in _table_bboxes:
            if r.overlaps(tb, threshold=0.3):
                return True
        return False

    # Collect bboxes of non-table drawings
    draw_rects: list[BBox] = []
    for d in drawings:
        r = d.get("rect")
        if r is None:
            continue
        rect = (r.x0, r.y0, r.x1, r.y1)
        if _in_table(rect):
            continue
        bbox = BBox(r.x0, r.y0, r.x1, r.y1)
        # Skip very thin lines (likely rules/separators, not diagram elements)
        if bbox.width < 2 and bbox.height < 2:
            continue
        draw_rects.append(bbox)

    if len(draw_rects) < 10:
        return []

    # Cluster drawings spatially using union-find + spatial grid
    # Grid-based approach: O(n) average instead of O(n²)
    n = len(draw_rects)
    parent = list(range(n))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    gap = 15.0  # allow slightly larger gap for vector elements
    cell_size = max(gap * 2, 30.0)  # grid cell size

    # Build spatial grid: map (cell_x, cell_y) -> list of rect indices
    grid: dict[tuple[int, int], list[int]] = {}
    for i, b in enumerate(draw_rects):
        # Each rect may span multiple grid cells
        cx0 = int((b.x0 - gap) / cell_size)
        cy0 = int((b.y0 - gap) / cell_size)
        cx1 = int((b.x1 + gap) / cell_size)
        cy1 = int((b.y1 + gap) / cell_size)
        for cx in range(cx0, cx1 + 1):
            for cy in range(cy0, cy1 + 1):
                grid.setdefault((cx, cy), []).append(i)

    # Check overlaps only within same grid cells
    for cell_indices in grid.values():
        for ii in range(len(cell_indices)):
            for jj in range(ii + 1, len(cell_indices)):
                i, j = cell_indices[ii], cell_indices[jj]
                bi, bj = draw_rects[i], draw_rects[j]
                if (bi.x0 - gap <= bj.x1 and bj.x0 - gap <= bi.x1
                        and bi.y0 - gap <= bj.y1 and bj.y0 - gap <= bi.y1):
                    union(i, j)

    # Group by cluster root
    clusters: dict[int, list[int]] = {}
    for i in range(n):
        root = find(i)
        clusters.setdefault(root, []).append(i)

    figures: list[Figure] = []
    for indices in clusters.values():
        # Require minimum complexity: ≥10 drawing commands
        if len(indices) < 10:
            continue
        # Compute bounding box of cluster
        x0 = min(draw_rects[i].x0 for i in indices)
        y0 = min(draw_rects[i].y0 for i in indices)
        x1 = max(draw_rects[i].x1 for i in indices)
        y1 = max(draw_rects[i].y1 for i in indices)
        # Require minimum area (50x50 pt)
        if (x1 - x0) < 50 or (y1 - y0) < 50:
            continue
        figures.append(Figure(
            image_index=-1,  # no raster image index
            bbox=BBox(x0, y0, x1, y1),
            page_num=page_num,
            is_vector=True,
        ))

    return figures


def extract_images(page: fitz.Page, page_num: int) -> list[Figure]:
    """Extract image regions from a page.

    Applies size filtering and spatial clustering to merge grid tiles
    (e.g. generated samples in diffusion model papers) into single figures.
    """
    raw_figures: list[Figure] = []
    image_list = page.get_images(full=True)
    for idx, img in enumerate(image_list):
        xref = img[0]
        # Get image bbox from page
        rects = page.get_image_rects(xref)
        if rects:
            rect = rects[0]
            bbox = BBox(rect.x0, rect.y0, rect.x1, rect.y1)
            # Skip tiny images (icons, decorators, bullets)
            if bbox.width < 30 or bbox.height < 30:
                continue
            raw_figures.append(Figure(
                image_index=idx,
                bbox=bbox,
                page_num=page_num,
                xref=xref,
            ))

    if not raw_figures:
        return []

    # Spatial clustering: merge images that are close together into one figure.
    # This handles grids of generated samples common in generative model papers.
    return _cluster_nearby_figures(raw_figures)


def _cluster_nearby_figures(figures: list[Figure], gap: float = 10.0) -> list[Figure]:
    """Merge figures whose bounding boxes are within *gap* points of each other.

    Uses single-linkage clustering: if any two images in a cluster are close,
    they belong to the same cluster.  The merged figure uses the bounding-box
    union and keeps the first image_index.
    """
    if len(figures) <= 1:
        return figures

    n = len(figures)
    parent = list(range(n))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    # O(n^2) but n is small per-page
    for i in range(n):
        for j in range(i + 1, n):
            bi, bj = figures[i].bbox, figures[j].bbox
            # Check if bboxes are within gap points (expand each bbox by gap)
            if (bi.x0 - gap <= bj.x1 and bj.x0 - gap <= bi.x1
                    and bi.y0 - gap <= bj.y1 and bj.y0 - gap <= bi.y1):
                union(i, j)

    # Group by cluster root
    clusters: dict[int, list[int]] = {}
    for i in range(n):
        root = find(i)
        clusters.setdefault(root, []).append(i)

    merged: list[Figure] = []
    for indices in clusters.values():
        figs = [figures[i] for i in indices]
        # Merge bboxes
        x0 = min(f.bbox.x0 for f in figs)
        y0 = min(f.bbox.y0 for f in figs)
        x1 = max(f.bbox.x1 for f in figs)
        y1 = max(f.bbox.y1 for f in figs)
        # When multiple tiles are merged, set xref=None so the merged
        # bbox will be rendered as a composite image in pdf.py (single
        # xref would only capture the first tile).
        merged.append(Figure(
            image_index=figs[0].image_index,
            bbox=BBox(x0, y0, x1, y1),
            page_num=figs[0].page_num,
            caption=figs[0].caption,
            xref=figs[0].xref if len(figs) == 1 else None,
        ))
    return merged


def find_captions(lines: list[TextLine], page_num: int) -> list[DocumentBlock]:
    """Find figure and table captions in text lines.

    Detects captions that start with:
      - "Figure 1: ..."  / "Figure 1. ..."  / "Figure 1 ..."
      - "Fig. 1 ..."     / "Fig 1. ..."
      - "Table 1: ..."   / "Table I." ...  (Roman numerals)
      - "TABLE I. ..."   (all caps)
      - "Table 3.1: ..."  (compound numbers)

    Handles multi-line captions where continuation lines are at a similar
    horizontal position and vertically close.  Captions in a smaller font
    than surrounding body text are also detected.
    """
    captions: list[DocumentBlock] = []

    i = 0
    while i < len(lines):
        text = lines[i].text.strip()
        m = _match_caption(text)
        if m:
            # Collect continuation lines
            caption_lines = [lines[i]]
            ref_x0 = lines[i].bbox.x0  # anchor x for continuation detection
            j = i + 1
            while j < len(lines):
                next_line = lines[j]
                next_text = next_line.text.strip()
                # Stop on blank line
                if not next_text:
                    break
                # Stop if next line is itself a caption start
                if _match_caption(next_text):
                    break
                y_gap = next_line.bbox.y0 - caption_lines[-1].bbox.y1
                x_diff = abs(next_line.bbox.x0 - ref_x0)
                # Continuation: close vertically (within 18pt) and similar x
                # (allow up to 40pt x-shift for indented continuation lines)
                if y_gap <= 18 and x_diff <= 40:
                    caption_lines.append(next_line)
                    j += 1
                else:
                    break

            full_text = " ".join(l.text.strip() for l in caption_lines)
            from ..utils.geometry import merge_bboxes
            bbox = merge_bboxes([l.bbox for l in caption_lines])
            num_str = m.group("number")
            cap_number = _parse_caption_number(num_str)
            captions.append(DocumentBlock(
                block_type=BlockType.CAPTION,
                text=full_text,
                bbox=bbox,
                page_num=page_num,
                lines=caption_lines,
                metadata={
                    "type": _caption_type(m.group("keyword")),
                    "number": cap_number,
                    "number_str": num_str,
                },
            ))
            i = j
        else:
            i += 1

    return captions


def associate_captions(
    figures: list[Figure],
    captions: list[DocumentBlock],
) -> list[Figure]:
    """Associate captions with their nearest figure on the same page.

    Looks both above AND below the figure bounding box.  A caption is
    accepted if it is within 150pt vertically and horizontally overlaps
    the figure by at least a small amount.  When multiple captions are
    candidates, we prefer the closest one.
    """
    if not figures or not captions:
        return figures

    figure_captions = [
        c for c in captions
        if c.metadata.get("type") == "figure"
    ]

    for fig in figures:
        if not fig.bbox:
            continue
        best_caption = None
        best_dist = float("inf")
        for cap in figure_captions:
            # Must be on the same page
            if cap.page_num != fig.page_num:
                continue
            dist = vertical_distance(fig.bbox, cap.bbox)
            # Increase proximity threshold to 150pt to catch wide gaps in
            # multi-column layouts and tall figures
            if dist >= 150:
                continue
            # Prefer captions that horizontally overlap the figure
            # (use a lenient threshold — even partial overlap is fine)
            h_ov = horizontal_overlap(fig.bbox, cap.bbox)
            # Score: distance, penalised slightly when there is no overlap
            score = dist + (0 if h_ov > 0.05 else 20)
            if score < best_dist:
                best_dist = score
                best_caption = cap

        if best_caption is not None:
            fig.caption = best_caption.text

    return figures


def figure_to_markdown(fig: Figure) -> str:
    """Convert a Figure to Markdown representation."""
    caption = fig.caption or f"Figure (page {fig.page_num + 1})"
    # Clean caption for markdown
    caption = caption.strip()
    return f"![{caption}](figure_{fig.page_num + 1}_{fig.image_index + 1}.png)\n\n*{caption}*"
