"""Step 4: Structural pre-filter using natural-pdf (local, no API calls)."""

import logging
from typing import Dict, List, Tuple

from .config import (
    CHECKBOX_MAX_ASPECT_RATIO,
    CHECKBOX_MAX_SIZE,
    CHECKBOX_MIN_SIZE,
    CHECKBOX_SIGNALS,
    MIN_STRUCTURAL_SCORE,
)
from .download import get_pdf_path
from .prefilter import score_page
from .state import PipelineState

logger = logging.getLogger(__name__)


def is_checkbox_shaped(bbox: Tuple[float, float, float, float]) -> bool:
    """Check if a bounding box matches checkbox geometry.

    Uses the same criteria as VectorCheckboxOptions defaults:
    size 6-25 PDF points, aspect ratio <= 1.5.
    """
    x0, y0, x1, y1 = bbox
    w = x1 - x0
    h = y1 - y0

    if w <= 0 or h <= 0:
        return False
    if w < CHECKBOX_MIN_SIZE or h < CHECKBOX_MIN_SIZE:
        return False
    if w > CHECKBOX_MAX_SIZE or h > CHECKBOX_MAX_SIZE:
        return False

    aspect = max(w, h) / min(w, h)
    if aspect > CHECKBOX_MAX_ASPECT_RATIO:
        return False

    return True


def structural_filter(state: PipelineState) -> Dict[str, List[int]]:
    """Score candidate pages using natural-pdf structural analysis.

    Returns dict mapping doc_id -> list of qualifying page indices.
    """
    from natural_pdf import PDF

    results: Dict[str, List[int]] = {}
    processed = 0
    skipped = 0

    for doc_id in state.all_doc_ids():
        # Must be downloaded first
        if not state.is_done(doc_id, "downloaded"):
            continue

        # Skip if already done
        if state.is_done(doc_id, "structural_filtered"):
            cached = state.get(doc_id, "qualifying_pages", [])
            if cached:
                results[doc_id] = cached
            continue

        candidate_pages = state.get(doc_id, "candidate_pages", [])
        if not candidate_pages:
            state.set_batch(doc_id, {"structural_filtered": True, "qualifying_pages": []})
            skipped += 1
            continue

        pdf_path = get_pdf_path(doc_id)
        if not pdf_path.exists():
            logger.warning("PDF not found for doc %s", doc_id)
            skipped += 1
            continue

        qualifying_pages = []
        try:
            pdf = PDF(str(pdf_path))
            try:
                for page_idx in candidate_pages:
                    if page_idx >= len(pdf.pages):
                        continue

                    page = pdf.pages[page_idx]
                    score = _score_page_structural(page)

                    if score >= MIN_STRUCTURAL_SCORE:
                        qualifying_pages.append(page_idx)
            finally:
                pdf.close()
        except Exception as e:
            logger.warning("Failed to process doc %s: %s", doc_id, e)
            state.set_batch(doc_id, {"structural_filtered": True, "qualifying_pages": []})
            skipped += 1
            continue

        state.set_batch(doc_id, {
            "structural_filtered": True,
            "qualifying_pages": qualifying_pages,
        })

        if qualifying_pages:
            results[doc_id] = qualifying_pages

        processed += 1
        if processed % 50 == 0:
            logger.info("Structural filter: processed %d documents...", processed)

    total_pages = sum(len(p) for p in results.values())
    logger.info(
        "Structural filter complete. %d docs with %d qualifying pages. %d skipped.",
        len(results), total_pages, skipped,
    )
    return results


def _score_page_structural(page) -> int:
    """Score a single page for checkbox likelihood using structural signals."""
    score = 0

    # 1. Small square rects (vector checkbox candidates)
    try:
        rects = page.find_all("rect")
        checkbox_rects = [r for r in rects if is_checkbox_shaped(r.bbox)]
        score += len(checkbox_rects)
    except Exception:
        pass

    # 2. AcroForm widget annotations (fillable form fields)
    try:
        raw_annots = page._page.annots or []
        widget_btns = [a for a in raw_annots if "/Btn" in str(a.get("FT", ""))]
        score += len(widget_btns) * 2
    except Exception:
        pass

    # 3. Text signal confirmation
    try:
        text = page.extract_text() or ""
        score += score_page(text)
    except Exception:
        pass

    return score


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    st = PipelineState()
    structural_filter(st)
