"""Cross-page article-continuation linker (classical, no Claude / no API).

Telugu papers "jump" long stories across pages:

  • bottom of the source article :  'మిగతా <N>వ పేజీలో...'   (rest is on page N)
  • top of the continuation page  :  the SAME headline repeated, plus the marker
                                     '(మొదటి పేజీ తరువాయి)'   (continued from page 1)

We REQUIRE BOTH signals: the jump-from marker on the source page (which also tells
us the target page N) AND a repeated headline on page N that matches the source
article's headline. We then build a boundary box for the continuation by injecting
a SYNTHETIC dateline anchored on that repeated headline and running the existing
boundary engine (so the continuation floors correctly against page-N's own
articles), and finally STITCH the source crop + continuation crop into ONE tall
image per article.

Pure geometry + Tesseract on saved fixtures (page PNGs + regions.json + the
per-article info.json the pipeline already wrote). No Claude.

    link_continuations(run_dir) -> list of link records (also written to
        <run>/continuations/links.json, stitched PNGs alongside).
"""
import difflib
import json
import re
import unicodedata
from pathlib import Path

from PIL import Image

from _lines import detect_separator_lines

try:
    from smart_extractor import _dateline_in_text
except Exception:                       # keep the module importable in isolation
    def _dateline_in_text(_t, _paper=None):
        return None

_ZW = "‌‍"                       # zero-width non-joiner / joiner
_TERMINATORS = ".।?!"            # sentence-enders: full stop, danda, ?, !
_TELUGU_DIGITS = {ord("౦") + i: str(i) for i in range(10)}

# jump-FROM: 'మిగతా <num> వ పేజీ...' — tolerate OCR noise between the tokens.
_JUMP_FROM = re.compile(r"మిగ[తథధ]ా.{0,10}?([0-9౦-౯]{1,2})\s*వ?\s*పే[జీిౌ]", re.UNICODE)
# jump-TO: 'మొదటి పేజీ తరువాయి' — the continuation marker on the target page.
_JUMP_TO = re.compile(r"మొదట[ిి]?\s*పే[జీి].{0,6}?తరు[వ]?ా[యి]", re.UNICODE)


def _paper_of_height(h):
    if 7040 <= h <= 7060:
        return "sakshi"
    if h in (6833, 6549):
        return "andhra_jyothi"
    if h == 6422:
        return "namaste_telangana"
    return f"h{h}"


def _ocr(page, box):
    import pytesseract
    try:
        return pytesseract.image_to_string(page.crop((int(box[0]), int(box[1]),
                                                       int(box[2]), int(box[3]))),
                                            lang="tel")
    except Exception:
        return ""


def _page_num_from(text):
    m = _JUMP_FROM.search(text or "")
    if not m:
        return None
    digits = m.group(1).translate(_TELUGU_DIGITS)
    dm = re.search(r"\d+", digits)
    if not dm:
        return None
    d = dm.group()
    # OCR frequently DOUBLE-RECOGNISES a single glyph as the Latin digit followed
    # by its Telugu twin (e.g. '3౩' → '33'). District editions only ever run a
    # handful of pages, so a run of identical digits is that artefact, not page 33.
    if len(d) == 2 and d[0] == d[1]:
        d = d[0]
    return int(d)


def _norm(s):
    s = unicodedata.normalize("NFC", s or "")
    s = "".join(c for c in s if c not in _ZW)
    return re.sub(r"[\s.,!?:;\-()।\"'‘’]", "", s)


def _similar(a, b):
    a, b = _norm(a), _norm(b)
    if not a or not b:
        return 0.0
    return difflib.SequenceMatcher(None, a, b).ratio()


def _find_jump_from(page, art_bbox):
    """Scan the BOTTOM band of the source article for the 'మిగతా Nవ పేజీలో' marker;
    fall back to the whole article box. Returns the target page number or None."""
    l, t, r, b = art_bbox
    band_top = max(t, b - 340)
    n = _page_num_from(_ocr(page, (l, band_top, r, b)))
    if n is None:                                   # marker may sit a little higher
        n = _page_num_from(_ocr(page, (l, t, r, b)))
    return n


def _best_headline(regs, page, headline_tel, thresh=0.60):
    """Among the doc_titles on the target page, return (region, score, has_marker)
    for the one whose OCR text best matches the source headline (>= thresh)."""
    best, best_score = None, 0.0
    for r in regs:
        if r.get("type") != "doc_title":
            continue
        score = _similar(_ocr(page, r["bbox"]), headline_tel)
        if score > best_score:
            best, best_score = r, score
    if not best or best_score < thresh:
        return None, best_score, False
    hb = best["bbox"]
    below = _ocr(page, (hb[0], hb[3], hb[2], min(hb[3] + 240, page.height)))
    has_marker = bool(_JUMP_TO.search(below) or _JUMP_TO.search(_ocr(page, hb)))
    return best, best_score, has_marker


def _ends_terminated(text):
    """True if the column's last word ends a sentence (., danda, ?, !), ignoring
    trailing quotes/brackets. This is the (noisy, OCR-based) 'article concluded'
    signal — used only AFTER the structural stop-tests have had their say."""
    t = (text or "").rstrip()
    t = t.rstrip("\"'’”)］]》」』 ")
    return bool(t) and t[-1] in _TERMINATORS


def _col_width(regs):
    """Robust single-column width = median width of the page's text boxes."""
    ws = sorted(r["bbox"][2] - r["bbox"][0] for r in regs if r.get("type") == "text")
    return ws[len(ws) // 2] if ws else 700


def _has_body_below(regs, b):
    """True if a text box sits below box `b` within its horizontal span — i.e. `b`
    is a headline opening its OWN article (a new story), not a trailing label or an
    inline emphasis run. Used to let a paragraph_title act as an article floor."""
    for r in regs:
        if r.get("type") == "text":
            rb = r["bbox"]
            if rb[1] > b[1] + 5 and rb[0] < b[2] and rb[2] > b[0]:
                return True
    return False


def _is_floor_head(regs, r, hid):
    """A region that floors a continuation: any doc_title, OR a paragraph_title that
    has its own body beneath it (the start of a DIFFERENT article — e.g. a
    'continued from page N' sub-head sitting below our continuation block)."""
    if r.get("id") == hid:
        return False
    t = r.get("type")
    if t == "doc_title":
        return True
    return t == "paragraph_title" and _has_body_below(regs, r["bbox"])


def _floor_under(regs, lines, hb, hid, x0, x1, H):
    """Lowest boundary below the header within the corridor [x0,x1]: the top of the
    next headline (doc_title, or an article-starting paragraph_title), or a
    horizontal rule, whichever comes first (else page bottom)."""
    f = H
    for r in regs:
        b = r["bbox"]
        if (_is_floor_head(regs, r, hid)
                and b[1] > hb[3] + 20 and b[0] < x1 and b[2] > x0):
            f = min(f, b[1])
    for ln in (lines.get("h") or []):
        y = ln["y1"]
        if hb[3] + 20 < y < f and ln["x1"] < x1 and ln["x2"] > x0:
            f = min(f, y)
    return f


def _floor_below_body(regs, lines, hb, hid, x0, x1, body, H):
    """Lowest article boundary in corridor [x0,x1]. A next-headline or horizontal
    rule only counts as the floor when NO included body text continues *below* it in
    an overlapping column — that distinguishes the next article's wide headline (no
    body beneath, within this block) from an interior sub-head / column-top rule
    (article body keeps going underneath it)."""
    def _no_body_below(bx0, bx2, yt):
        for bb in body:
            if bb[3] > yt + 10 and bb[0] < bx2 and bb[2] > bx0:
                return False
        return True

    f = H
    for r in regs:
        b = r["bbox"]
        if (_is_floor_head(regs, r, hid)
                and b[1] > hb[3] + 20 and b[0] < x1 and b[2] > x0
                and _no_body_below(b[0], b[2], b[1])):
            f = min(f, b[1])
    for ln in (lines.get("h") or []):
        y = ln["y1"]
        if hb[3] + 20 < y < f and ln["x1"] < x1 and ln["x2"] > x0 \
                and _no_body_below(ln["x1"], ln["x2"], y):
            f = min(f, y)
    return f


def _vrule_in_gutter(lines, x0, x1, y0, y1):
    """A printed vertical rule sitting in the gutter [x0,x1] across most of [y0,y1]."""
    for ln in (lines.get("v") or []):
        if x0 - 8 <= ln["x1"] <= x1 + 8:
            ov = min(ln["y2"], y1) - max(ln["y1"], y0)
            if ov > 0.4 * (y1 - y0):
                return True
    return False


def _track_starts_with_title(regs, items, hb):
    """STOP-test 5: does a headline / sub-head box open the next column?"""
    xs = min(it["bbox"][0] for it in items)
    xe = max(it["bbox"][2] for it in items)
    ytop = min(it["bbox"][1] for it in items)
    for r in regs:
        if r.get("type") in ("doc_title", "paragraph_title"):
            b = r["bbox"]
            if b[0] < xe and b[2] > xs and b[1] > hb[3] + 20 and abs(b[1] - ytop) < 130:
                return True
    return False


def _continuation_crop(regs, png_path, paper, headline_region):
    """Crop the FULL continuation article by GROWING RIGHT from the headline.

    A continuation's body is usually wider than its header and spans several columns.
    We start at the header's column and walk right column-by-column, absorbing the
    next column unless a STOP-test fires:

      1 no next column            5 next column opens with a headline/sub-head
      2 vertical-rule wall        6 current column's last word ends with a period
      3 horizontal-rule break     7 grammatical bridge fails (approx: 4+6 absent)
      4 next column opens a dateline

    Structural tests (1-5) outrank the OCR period (6). The crop is the single
    rectangle enclosing every absorbed column (+ any photo inside the block),
    floored at the next headline / horizontal rule.

    Returns (PIL.Image, debug_bbox) or (None, None)."""
    page = Image.open(png_path).convert("RGB")
    H = page.height
    hb = headline_region["bbox"]
    hid = headline_region.get("id")
    cx = (hb[0] + hb[2]) / 2
    lines = detect_separator_lines(png_path)
    colW = _col_width(regs)

    # provisional floor under the header's own span gives us the band to cluster in
    F0 = _floor_under(regs, lines, hb, hid, hb[0], hb[2], H)

    # body text boxes in the band, clustered into left-to-right COLUMN TRACKS
    band = [r for r in regs if r.get("type") == "text"
            and (r["bbox"][1] + r["bbox"][3]) / 2 < F0
            and r["bbox"][3] > hb[1]
            and r["bbox"][0] >= hb[0] - colW * 0.5]   # drop columns left of the header
    if not band:
        return None, None
    tracks = []
    for r in sorted(band, key=lambda r: r["bbox"][0]):
        for tr in tracks:
            if abs(r["bbox"][0] - tr["x0"]) < colW * 0.5:
                tr["items"].append(r)
                break
        else:
            tracks.append({"x0": r["bbox"][0], "items": [r]})
    tracks.sort(key=lambda t: t["x0"])

    # START column = the track straddling the headline centre (else nearest)
    start = None
    for i, t in enumerate(tracks):
        if min(it["bbox"][0] for it in t["items"]) <= cx <= max(it["bbox"][2] for it in t["items"]):
            start = i
            break
    if start is None:
        start = min(range(len(tracks)), key=lambda i: abs(tracks[i]["x0"] - cx))

    import os
    DBG = os.environ.get("CONT_DEBUG")
    if DBG:
        print(f"[dbg] hb={hb} cx={cx} colW={colW} F0={F0}")
        for ti, t in enumerate(tracks):
            print(f"[dbg] track {ti} x0={t['x0']} items="
                  + ",".join(f"{it.get('id')}{it['bbox']}" for it in t['items']))
        print(f"[dbg] start track = {start}")

    # GROW RIGHT, applying the stop-tests at each gutter
    included = [start]
    i = start
    while True:
        cur = tracks[i]["items"]
        cur_r = max(it["bbox"][2] for it in cur)
        nxt = next((j for j in range(i + 1, len(tracks))
                    if tracks[j]["x0"] > cur_r - colW * 0.3), None)
        if nxt is None:                                         # 1 no next column
            if DBG: print(f"[dbg] from {i}: STOP test1 no next column")
            break
        nt = tracks[nxt]["items"]
        nt_l = min(it["bbox"][0] for it in nt)
        y0, y1 = hb[1], F0
        if _vrule_in_gutter(lines, cur_r, nt_l, y0, y1):        # 2 vertical wall
            if DBG: print(f"[dbg] from {i} to {nxt}: STOP test2 vrule wall {cur_r}-{nt_l}")
            break
        top_reg = min(nt, key=lambda it: it["bbox"][1])
        tb = top_reg["bbox"]
        head_txt = _ocr(page, (tb[0], tb[1], tb[2], min(tb[1] + 150, tb[3])))
        if _dateline_in_text(head_txt, paper):                  # 4 dateline opens next col
            if DBG: print(f"[dbg] from {i} to {nxt}: STOP test4 dateline '{head_txt[:40]}'")
            break
        if _track_starts_with_title(regs, nt, hb):              # 5 headline opens next col
            if DBG: print(f"[dbg] from {i} to {nxt}: STOP test5 title opens next")
            break
        bot_reg = max(cur, key=lambda it: it["bbox"][3])
        bot_txt = _ocr(page, bot_reg["bbox"])
        if _ends_terminated(bot_txt):                           # 6 period concludes
            if DBG: print(f"[dbg] from {i} to {nxt}: STOP test6 period; tail='{bot_txt[-30:]}'")
            break
        if DBG: print(f"[dbg] grow {i} -> {nxt}")
        included.append(nxt)                                    # 7 bridge holds -> grow
        i = nxt
    if DBG: print(f"[dbg] included tracks = {included}")

    # FINAL rectangle over the absorbed columns, floored against the real corridor
    left = hb[0]
    right = max(max(it["bbox"][2] for it in tracks[k]["items"]) for k in included)
    body = [it["bbox"] for k in included for it in tracks[k]["items"]]
    F = _floor_below_body(regs, lines, hb, hid, left, right, body, H)
    if DBG: print(f"[dbg] left={left} right={right} F={F}")

    members = list(body)
    for r in regs:                       # sweep in photos / captions inside the block
        if r.get("type") in ("image", "figure_title", "paragraph_title", "doc_title"):
            b = r["bbox"]
            bcx = (b[0] + b[2]) / 2
            if left - 10 <= bcx <= right + 10 and b[3] <= F and b[3] > hb[1] - 320:
                members.append(b)

    left = min(left, min(m[0] for m in members))
    top = min(hb[1], min(m[1] for m in members))
    bot = min(int(F) - 8, max(m[3] for m in members) + 22)
    box = (int(left), int(max(0, top - 6)), int(right), int(bot))
    return page.crop(box), list(box)


def _stitch(top_png, cont_img, out_path):
    """Stack the source page-1 crop above the composed continuation image."""
    top = Image.open(top_png).convert("RGB")
    W = max(top.width, cont_img.width)
    sep = 18
    combo = Image.new("RGB", (W, top.height + sep + cont_img.height), "white")
    combo.paste(top, (0, 0))
    combo.paste(cont_img, (0, top.height + sep))
    combo.save(out_path)
    return combo.size


def link_continuations(run_dir, out_dir=None, match_thresh=0.60, verbose=True):
    run = Path(run_dir)
    pages_dir = run / "pages"
    arts_dir = run / "articles"
    out_dir = Path(out_dir) if out_dir else (run / "continuations")
    out_dir.mkdir(parents=True, exist_ok=True)

    p1 = pages_dir / "page_001.png"
    paper = _paper_of_height(Image.open(p1).height) if p1.exists() else "unknown"

    results = []
    for info_path in sorted(arts_dir.glob("*/info.json")):
        info = json.loads(info_path.read_text())
        P, bbox = info.get("page"), info.get("bbox")
        headline = info.get("headline_telugu") or ""
        if not P or not bbox:
            continue
        # Continuations only ever START on page 1 — the front page jumps long lead
        # stories onto inside pages. Articles on pages 2+ are self-contained, so we
        # never scan them for a jump-from marker.
        if P != 1:
            continue
        src_png = pages_dir / f"page_{P:03d}.png"
        if not src_png.exists():
            continue

        N = _find_jump_from(Image.open(src_png), bbox)        # SIGNAL 1: jump-from
        if not N or N == P:
            continue

        # The jump-from page DIGIT is OCR-noisy (Telugu ౪/౬ confusion) and can point
        # at a page outside this extraction (here the front page said '6వ పేజీ' but
        # the run only has 4 pages). So treat the marker as "this front-page lead
        # continues on a later page" and resolve the ACTUAL target by the repeated-
        # headline signal: score every OTHER page's doc_titles against the source
        # headline and take the best match, breaking ties toward the OCR'd page N.
        existing = sorted(
            int(m.group(1))
            for q in pages_dir.glob("page_*.png")
            if (m := re.match(r"page_(\d+)\.png$", q.name)))
        best = None   # (score, page, tregs, tgt_png_str, hl, marker)
        for q in existing:
            if q == P:
                continue
            qr = pages_dir / f"page_{q:03d}.regions.json"
            qp = pages_dir / f"page_{q:03d}.png"
            if not qr.exists() or not qp.exists():
                continue
            qregs = json.loads(qr.read_text())["regions"]
            hl_q, score_q, marker_q = _best_headline(qregs, Image.open(qp), headline, match_thresh)
            if hl_q is None:
                continue
            if best is None or score_q > best[0] or (score_q == best[0] and q == N):
                best = (score_q, q, qregs, str(qp), hl_q, marker_q)

        if best is None:
            rec = {"article": info_path.parent.name, "src_page": P, "headline": headline,
                   "target_page": N, "match_score": 0.0, "taruvaayi_marker": False,
                   "linked": False,
                   "reason": "no matching repeated headline on any inside page"}
            results.append(rec)
            if verbose:
                print(f"  ✗ {rec['article']} → p{N}: jump-from found but no repeated "
                      f"headline match on any inside page (NOT linked)")
            continue

        score, N, tregs, tgt_png, hl, marker = best
        rec = {"article": info_path.parent.name, "src_page": P, "headline": headline,
               "target_page": N, "match_score": round(score, 2),
               "taruvaayi_marker": marker, "linked": False}

        cont_img, cbox = _continuation_crop(tregs, tgt_png, paper, hl)
        if cont_img is None:
            rec["reason"] = "no body under repeated headline"
            results.append(rec)
            continue
        out_png = out_dir / f"{info_path.parent.name}__cont_p{N:03d}.png"
        size = _stitch(arts_dir / info_path.parent.name / "article.png",
                       cont_img, out_png)
        rec.update({"linked": True, "cont_bbox": cbox,
                    "stitched": str(out_png), "stitched_size": size})
        results.append(rec)
        if verbose:
            print(f"  ✓ {rec['article']} '{headline[:26]}' → p{N} "
                  f"(headline {score:.2f}, తరువాయి={marker}) → {out_png.name}")

    (out_dir / "links.json").write_text(json.dumps(results, indent=2, ensure_ascii=False))
    linked = sum(1 for r in results if r["linked"])
    if verbose:
        print(f"\n  {linked} continuation(s) linked & stitched; report: {out_dir/'links.json'}")
    return results


if __name__ == "__main__":
    import sys
    link_continuations(sys.argv[1] if len(sys.argv) > 1 else "output/20260601_234709")