detect.py

"""
Core PDF field detection logic.
  1. AcroForm PDF      → extract native widgets
  2. Flat vector PDF   → extract from drawings layer (rectangles, lines)
  3. Image/scanned PDF → run FFDNet (commonforms) for checkbox/text detection
"""
import fitz
import tempfile, os
from typing import Literal

# FFDNet via commonforms — loaded lazily on first use
_ffdnet_ready = False
_ffdnet_error = ''
def _ensure_ffdnet():
    global _ffdnet_ready, _ffdnet_error
    if not _ffdnet_ready and not _ffdnet_error:
        try:
            import commonforms as _cf
            _ffdnet_ready = True
        except Exception as e:
            import traceback
            _ffdnet_error = traceback.format_exc()
            print(f"commonforms import failed:\n{_ffdnet_error}")
    return _ffdnet_ready


FieldType = Literal['checkbox', 'text', 'signature']


def _label_near(x0, y0, x1, y1, text_spans):
    """Find the nearest text label above or left of a rect."""
    best, best_dist = '', 1e9
    for span in text_spans:
        sx0, sy0, sx1, sy1 = span['bbox']
        # Candidate: text ends to the left, or is above (within 20pt)
        if sx1 <= x0 + 5 and abs((sy0 + sy1) / 2 - (y0 + y1) / 2) < 20:
            dist = x0 - sx1
            if 0 <= dist < best_dist:
                best, best_dist = span['text'], dist
        elif sy1 <= y0 + 2 and sx0 >= x0 - 5 and sx1 <= x1 + 5:
            dist = y0 - sy1
            if 0 <= dist < best_dist:
                best, best_dist = span['text'], dist
    return best.strip()


def detect_page(page) -> dict:
    pw, ph = page.rect.width, page.rect.height
    widgets  = list(page.widgets())
    drawings = page.get_drawings()
    images   = page.get_images(full=False)

    # ── Case 1: AcroForm ────────────────────────────────────────────────────
    if widgets:
        boxes = []
        for w in widgets:
            r = w.rect
            ftype: FieldType
            if w.field_type in (fitz.PDF_WIDGET_TYPE_CHECKBOX, fitz.PDF_WIDGET_TYPE_RADIOBUTTON):
                ftype = 'checkbox'
            elif w.field_type == fitz.PDF_WIDGET_TYPE_SIGNATURE:
                ftype = 'signature'
            else:
                ftype = 'text'
            boxes.append({
                'type': ftype,
                'x': r.x0 / pw, 'y': (ph - r.y1) / ph,
                'w': r.width  / pw, 'h': r.height / ph,
                'label': w.field_name or '',
                'source': 'acroform',
            })
        return {'source': 'acroform', 'boxes': boxes}

    # ── Case 3: Image/scanned — run FFDNet ───────────────────────────────
    if not drawings:
        return {'source': 'needs_ml', 'boxes': []}  # resolved in detect_pdf

    # ── Case 2: Flat vector PDF ────────────────────────────────────────────
    text_spans = []
    for block in page.get_text('dict', flags=fitz.TEXT_INHIBIT_SPACES).get('blocks', []):
        for line in block.get('lines', []):
            for span in line.get('spans', []):
                t = span.get('text', '').strip()
                if t:
                    text_spans.append({'text': t, 'bbox': span['bbox']})

    boxes = []
    seen  = set()

    for d in drawings:
        r = d['rect']
        w, h = r.width, r.height
        if w < 1 or h < 1:
            continue

        key = (round(r.x0), round(r.y0))
        if key in seen:
            continue
        seen.add(key)

        x_frac = r.x0 / pw
        y_frac = r.y0 / ph
        w_frac = w / pw
        h_frac = h / ph

        # Small square → checkbox / radio
        if abs(w - h) < w * 0.35 and 3 < w < 22:
            label = _label_near(r.x0, r.y0, r.x1, r.y1, text_spans)
            boxes.append({
                'type': 'checkbox', 'source': 'vector',
                'x': x_frac, 'y': y_frac, 'w': w_frac, 'h': h_frac,
                'label': label,
            })

        # Thin horizontal line → text underline input
        elif h < 2.5 and w > 20:
            label = _label_near(r.x0, r.y0, r.x1, r.y1, text_spans)
            pad   = min(14 / ph, 0.02)
            boxes.append({
                'type': 'text', 'source': 'vector_line',
                'x': x_frac, 'y': max(0, y_frac - pad),
                'w': w_frac, 'h': pad + h_frac + 1 / ph,
                'label': label,
            })

        # Rectangular box wider than tall → text input field
        elif w > h * 1.2 and h > 6 and w < pw * 0.95:
            label = _label_near(r.x0, r.y0, r.x1, r.y1, text_spans)
            boxes.append({
                'type': 'text', 'source': 'vector_rect',
                'x': x_frac, 'y': y_frac, 'w': w_frac, 'h': h_frac,
                'label': label,
            })

    return {'source': 'vector', 'boxes': boxes}


def _run_ffdnet(pdf_bytes: bytes, page_nums: list[int]) -> dict[int, list[dict]]:
    """Run commonforms FFDNet on specific pages, return boxes per page index."""
    if not _ensure_ffdnet():
        return {}
    from commonforms import prepare_form

    with tempfile.TemporaryDirectory() as tmp:
        in_path  = os.path.join(tmp, 'in.pdf')
        out_path = os.path.join(tmp, 'out.pdf')
        with open(in_path, 'wb') as f:
            f.write(pdf_bytes)

        try:
            prepare_form(in_path, out_path, confidence=0.1, device='cpu')
            out_size = os.path.getsize(out_path) if os.path.exists(out_path) else 0
            print(f"FFDNet ran OK, output size={out_size} bytes")
        except Exception as e:
            import traceback
            print(f"FFDNet error: {e}")
            traceback.print_exc()
            return {}

        out_doc = fitz.open(out_path)
        results: dict[int, list[dict]] = {}
        for page_num in page_nums:
            if page_num >= len(out_doc):
                continue
            page = out_doc[page_num]
            pw, ph = page.rect.width, page.rect.height
            boxes = []
            for w in page.widgets():
                r = w.rect
                ftype = 'checkbox' if w.field_type in (
                    fitz.PDF_WIDGET_TYPE_CHECKBOX,
                    fitz.PDF_WIDGET_TYPE_RADIOBUTTON,
                ) else 'text'
                boxes.append({
                    'type': ftype,
                    'x': r.x0 / pw, 'y': r.y0 / ph,
                    'w': r.width  / pw, 'h': r.height / ph,
                    'label': w.field_name or '',
                    'source': 'ffdnet',
                })
            results[page_num] = boxes
        out_doc.close()
    return results


def detect_pdf(pdf_bytes: bytes) -> list[dict]:
    doc   = fitz.open(stream=pdf_bytes, filetype='pdf')
    pages = []

    for page_num, page in enumerate(doc):
        result = detect_page(page)
        result['page']   = page_num
        result['width']  = page.rect.width
        result['height'] = page.rect.height
        pages.append(result)

    doc.close()

    # Run FFDNet on any image pages in one pass (model loaded once)
    ml_pages = [p['page'] for p in pages if p['source'] == 'needs_ml']
    if ml_pages:
        ffdnet_results = _run_ffdnet(pdf_bytes, ml_pages)
        for p in pages:
            if p['source'] == 'needs_ml':
                p['source'] = 'ffdnet'
                p['boxes']  = ffdnet_results.get(p['page'], [])

    return pages