#!/usr/bin/env python3
"""
tables.py — selective table-cell extraction for the exam extractor (PLAN.md §B).

Two sources, unified into one cell-grid schema:
  * STANDARD  — the Tesseract+TableFormer backbone already emits `tables[].data.table_cells`
      (text + row/col offsets + spans + bbox). Free, cached, every run. Good on ruled tables;
      but it MISSES some data tables and OCRs them as loose tokens (REPORT.md p5).
  * GRANITE   — Granite-Docling-258M VLM emits `<otsl>` grids in DocTags (clean rows/cols even
      where the backbone scrambles them). GPU cost, so used SELECTIVELY: only on pages the router
      flags (a standard table present, or dense picture/checkbox), routed through dsync's GPU lock
      + Redis cache. Recipe (REPORT.md): {"to_formats":["doctags","json"], "pipeline":"vlm",
      "vlm_pipeline_model":"granite_docling"}.

Unified table = {page, n_rows, n_cols, grid (2D text), cells, caption, source, is_furniture}.
"""
import re, json, os, glob, base64, urllib.request

# ----------------------------------------------------------------- OTSL (Granite DocTags)
OTSL_BLOCK = re.compile(r"<otsl>(.*?)</otsl>", re.S)
CAPTION    = re.compile(r"<caption>(?:<loc_\d+>)*(.*?)</caption>", re.S)
CELL_TOK   = re.compile(r"<(fcel|ecel|ched|rhed|lcel|ucel|xcel|nl)>([^<]*)")
HEADER_TAGS = {"ched", "rhed"}


def parse_otsl(doctags):
    """Parse every <otsl> block in a DocTags string into unified tables."""
    out = []
    for block in OTSL_BLOCK.findall(doctags):
        cap = None
        mc = CAPTION.search(block)
        if mc:
            cap = re.sub(r"\s+", " ", mc.group(1)).strip()
        body = CAPTION.sub("", block)
        body = re.sub(r"<loc_\d+>", "", body)
        rows, cur = [], []
        for tag, txt in CELL_TOK.findall(body):
            if tag == "nl":
                rows.append(cur); cur = []
            else:
                cur.append({"text": txt.strip(), "header": tag in HEADER_TAGS,
                            "empty": tag == "ecel"})
        if cur:
            rows.append(cur)
        rows = [r for r in rows if r]
        if not rows:
            continue
        n_cols = max(len(r) for r in rows)
        grid = [[c["text"] for c in r] + [""] * (n_cols - len(r)) for r in rows]
        out.append({"page": None, "n_rows": len(rows), "n_cols": n_cols, "grid": grid,
                    "caption": cap, "source": "granite-otsl",
                    "is_furniture": is_furniture(grid, cap)})
    return out


# ----------------------------------------------------------------- standard TableFormer
def tables_from_standard(doc):
    out = []
    for t in doc.get("tables", []):
        data = t.get("data", {}) or {}
        cells = data.get("table_cells", []) or []
        nr, nc = data.get("num_rows") or 0, data.get("num_cols") or 0
        grid = [["" for _ in range(nc)] for _ in range(nr)]
        for c in cells:
            r0, c0 = c.get("start_row_offset_idx"), c.get("start_col_offset_idx")
            if r0 is not None and c0 is not None and r0 < nr and c0 < nc and c.get("text"):
                grid[r0][c0] = c["text"]
        prov = t.get("prov") or []
        page = prov[0].get("page_no") if prov else None
        cap = " ".join(x.get("text", "") for x in (t.get("captions") or []) if isinstance(x, dict)) or None
        out.append({"page": page, "n_rows": nr, "n_cols": nc, "grid": grid,
                    "caption": cap, "source": "docling-standard",
                    "is_furniture": is_furniture(grid, cap)})
    return out


# ----------------------------------------------------------------- furniture filter
FURNITURE_RE = re.compile(r"examiner|do not write|leave\s+blank|question\s*mark|"
                          r"for marker|total marks?$", re.I)


def is_furniture(grid, caption=None):
    """A table that is exam scaffolding (mark grid / 'For Examiner's Use'), not question data."""
    blob = " ".join(cell for row in grid for cell in row) + " " + (caption or "")
    if FURNITURE_RE.search(blob):
        return True
    # a single-column strip of question numbers / blanks = a mark grid
    flat = [c for row in grid for c in row if c.strip()]
    if flat and all(re.fullmatch(r"\d{1,2}", c.strip()) for c in flat):
        return True
    return False


# ----------------------------------------------------------------- Granite via dsync
VLM_OPTS = {"to_formats": ["doctags", "json"], "pipeline": "vlm",
            "vlm_pipeline_model": "granite_docling", "image_export_mode": "placeholder"}


def _serve_vlm(pdf_b64, fname, page):
    import dsync
    opts = {**VLM_OPTS, "page_range": [page, page]}
    body = {"options": opts,
            "sources": [{"kind": "file", "base64_string": pdf_b64, "filename": fname}],
            "target": {"kind": "inbody"}}
    req = urllib.request.Request(dsync.SERVE + "/v1/convert/source",
                                 data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    for _ in range(4):                                  # tolerate the single-use 404 race
        try:
            return json.loads(urllib.request.urlopen(req, timeout=1200).read())
        except urllib.error.HTTPError as e:
            if e.code == 404:
                import time; time.sleep(3); continue
            raise
    raise RuntimeError("serve vlm: repeated 404")


def _doctags_of(resp):
    doc = resp.get("document") or {}
    return doc.get("doctags_content") or doc.get("doc_tags") or doc.get("doctags") or ""


def granite_tables(pdf, pages, *, cached_glob=None, retries=4):
    """Run Granite-Docling on the given pages via dsync (GPU lock + OOM retry + Redis cache),
    parse <otsl>, tag each table with its page. Falls back to cached *.doctags if serve fails."""
    import dsync, time
    cache = _load_cached_doctags(cached_glob) if cached_glob else {}
    r = dsync._redis()
    b64 = base64.b64encode(open(pdf, "rb").read()).decode()
    fname = os.path.basename(pdf)
    sha = dsync._sha(pdf)
    out = []
    for pg in pages:
        key = f"docling:vlm:{sha}:p{pg}"
        doctags = None
        if r and (hit := r.get(key)):
            doctags = hit if isinstance(hit, str) else hit.decode()
        if doctags is None:
            delay = 5
            for attempt in range(retries):
                with dsync._GpuLock(r):
                    resp = _serve_vlm(b64, fname, pg)
                if dsync._is_oom(resp):
                    print(f"[granite] p{pg} OOM, backoff {delay}s ({attempt+1}/{retries})")
                    time.sleep(delay); delay = min(delay * 2, 120); continue
                doctags = _doctags_of(resp)
                if r and doctags:
                    r.set(key, doctags, ex=dsync.CACHE_TTL)
                break
        if not doctags and pg in cache:
            print(f"[granite] p{pg} serve empty -> cached doctags")
            doctags = cache[pg]
        for tbl in parse_otsl(doctags or ""):
            tbl["page"] = pg
            out.append(tbl)
    return out


def _load_cached_doctags(glob_pat):
    """Map page_no -> doctags text from files named *p<N>.doctags."""
    cache = {}
    for fn in glob.glob(glob_pat):
        m = re.search(r"p(\d+)\.doctags$", fn)
        if m:
            cache[int(m.group(1))] = open(fn, encoding="utf-8", errors="replace").read()
    return cache


# ----------------------------------------------------------------- routing + attach
def candidate_pages(doc):
    """Pages the router sends to Granite: a standard table, or a dense picture/checkbox page."""
    pages = set()
    for t in doc.get("tables", []):
        prov = t.get("prov") or []
        if prov and prov[0].get("page_no"):
            pages.add(prov[0]["page_no"])
    chk = {}
    for it in doc.get("texts", []):
        if it.get("label", "").startswith("checkbox"):
            prov = it.get("prov") or []
            if prov and prov[0].get("page_no"):
                chk[prov[0]["page_no"]] = chk.get(prov[0]["page_no"], 0) + 1
    pages |= {p for p, n in chk.items() if n >= 2}
    return sorted(pages)


def attach_to_questions(tables, parts):
    """Assign each non-furniture table to the nearest preceding part on its page (by y); if no
    geometry, attach to the first part on that page. Records table refs on the part."""
    data_tables = [t for t in tables if not t["is_furniture"]]
    by_page = {}
    for lab, v in parts.items():
        by_page.setdefault(v.get("page"), []).append((lab, v))
    for i, t in enumerate(data_tables):
        t["id"] = i
        cands = by_page.get(t["page"], [])
        if not cands:
            t["for_part"] = None; continue
        # best-effort: the part highest on the page (largest bbox top = the page's question stem),
        # else the earliest part label. (Tables sit under the stem; we don't carry table y here.)
        with_geo = [(lab, v) for lab, v in cands if v.get("bbox")]
        if with_geo:
            lab = max(with_geo, key=lambda kv: (kv[1]["bbox"] or {}).get("t", 0))[0]
        else:
            lab = sorted(cands, key=lambda kv: kv[0])[0][0]
        t["for_part"] = lab
        parts[lab].setdefault("tables", []).append(
            {"id": i, "n_rows": t["n_rows"], "n_cols": t["n_cols"],
             "caption": t["caption"], "source": t["source"]})
    return data_tables
