"""Tests for scripts/pdf_table_weave.py merge_tables.

CI-light: pure pikepdf fixtures built in-memory (+ saved/reopened for the
ParentTree-remap tests, matching test_table_to_figure.py's proven pattern),
no vendored ODL/java/verapdf.
"""

import os
import sys

import pikepdf

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "scripts"))

import pdf_table_weave as ptw  # noqa: E402


# ── fixtures ─────────────────────────────────────────────────────────────

def _cell(pdf, text=None, pg=None):
    kids = []
    if text is not None:
        p = pdf.make_indirect(pikepdf.Dictionary(
            {"/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/P"),
             "/ActualText": pikepdf.String(text)}))
        kids.append(p)
    d = {
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/TD"),
        "/K": pikepdf.Array(kids) if kids else pikepdf.Array([]),
    }
    if pg is not None:
        d["/Pg"] = pg
    return pdf.make_indirect(pikepdf.Dictionary(d))


def _row(pdf, cells):
    return pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/TR"),
        "/K": pikepdf.Array(cells),
    }))


def _group(pdf, tag, rows):
    """A /THead, /TBody, or /TFoot row-grouping struct elem — the PDF spec's
    standard tagged-table shape (/Table -> /TBody -> /TR -> /TD)."""
    return pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name(tag),
        "/K": pikepdf.Array(rows),
    }))


def _table(pdf, rows, caption=None, pg=None):
    kids = ([caption] if caption is not None else []) + list(rows)
    d = {
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Table"),
        "/K": pikepdf.Array(kids),
    }
    if pg is not None:
        d["/Pg"] = pg
    return pdf.make_indirect(pikepdf.Dictionary(d))


def _doc_with_tables(pdf, tables):
    doc = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Document"),
        "/K": pikepdf.Array(tables),
    }))
    for t in tables:
        t["/P"] = doc
    root = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructTreeRoot"), "/K": pikepdf.Array([doc]),
    }))
    pdf.Root.StructTreeRoot = root
    doc["/P"] = root
    return doc


def _walk_tags(node, acc):
    acc.append(str(node.get("/S", "")))
    for k in ptw._kids(node):
        _walk_tags(k, acc)


def _all_reachable_objgens(pdf):
    """objgen of every struct elem reachable from StructTreeRoot — used to
    assert a removed donor /Table is truly gone from the tree, not merely
    unlinked from one parent."""
    out = []

    def walk(n):
        if getattr(n, "is_indirect", False):
            out.append(n.objgen)
        for k in ptw._kids(n):
            walk(k)
    walk(pdf.Root.StructTreeRoot)
    return set(out)


def _two_frag_doc():
    """Two /Table fragments of the same 2-column shape — fragment 0 has the
    header row, fragment 1 is headerless (the real-world split-table shape
    described in the module docstring)."""
    pdf = pikepdf.new()
    t0 = _table(pdf, [
        _row(pdf, [_cell(pdf, "Name"), _cell(pdf, "Score")]),
        _row(pdf, [_cell(pdf, "Alice"), _cell(pdf, "9")]),
    ])
    t1 = _table(pdf, [
        _row(pdf, [_cell(pdf, "Bob"), _cell(pdf, "7")]),
        _row(pdf, [_cell(pdf, "Cara"), _cell(pdf, "8")]),
    ])
    _doc_with_tables(pdf, [t0, t1])
    return pdf


# ── basic merge: row count, order, objgen survival ─────────────────────────

def test_two_fragments_merge_row_count_sums_and_rows_survive_by_objgen():
    pdf = _two_frag_doc()
    tables = ptw._find_tables(pdf)
    t0_rows = ptw._table_rows(tables[0])
    t1_rows = ptw._table_rows(tables[1])
    original_keys = [r.objgen for r in t0_rows] + [r.objgen for r in t1_rows]

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])

    assert result == {
        "tables_merged": 1, "rows_moved": 2, "donors_removed": 1, "skipped": [],
    }

    tables_after = ptw._find_tables(pdf)
    assert len(tables_after) == 1
    survivor_rows = ptw._table_rows(tables_after[0])
    assert len(survivor_rows) == 4  # 2 + 2, summed
    assert [r.objgen for r in survivor_rows] == original_keys  # same objects, order preserved


def test_three_fragments_merge_in_given_order():
    pdf = pikepdf.new()
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "h1"), _cell(pdf, "h2")])])
    t1 = _table(pdf, [_row(pdf, [_cell(pdf, "a1"), _cell(pdf, "a2")])])
    t2 = _table(pdf, [_row(pdf, [_cell(pdf, "b1"), _cell(pdf, "b2")])])
    _doc_with_tables(pdf, [t0, t1, t2])

    result = ptw.merge_tables(pdf, [{"indices": [0, 1, 2]}])
    assert result["tables_merged"] == 1
    assert result["rows_moved"] == 2
    assert result["donors_removed"] == 2

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    rows = ptw._table_rows(tables[0])
    texts = [ptw._cell_text(ptw._row_cells(r)[0]) for r in rows]
    assert texts == ["h1", "a1", "b1"]  # survivor first, then donors in given order


# ── row-grouped fragments ───────────────────────────────────────────────────

def test_row_grouped_fragments_merge_and_leave_no_orphan_group_shells():
    pdf = pikepdf.new()
    thead = _group(pdf, "/THead", [_row(pdf, [_cell(pdf, "Name"), _cell(pdf, "Score")])])
    tbody0 = _group(pdf, "/TBody", [_row(pdf, [_cell(pdf, "Alice"), _cell(pdf, "9")])])
    t0 = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Table"),
        "/K": pikepdf.Array([thead, tbody0]),
    }))
    tbody1 = _group(pdf, "/TBody", [_row(pdf, [_cell(pdf, "Bob"), _cell(pdf, "7")])])
    t1 = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Table"),
        "/K": pikepdf.Array([tbody1]),
    }))
    _doc_with_tables(pdf, [t0, t1])
    tbody1_key = tbody1.objgen

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])
    assert result["tables_merged"] == 1
    assert result["rows_moved"] == 1
    assert result["skipped"] == []

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    rows = ptw._table_rows(tables[0])
    texts = [ptw._cell_text(ptw._row_cells(r)[0]) for r in rows]
    assert texts == ["Name", "Alice", "Bob"]  # THead row + survivor TBody row + donor row

    # donor's own /TBody row-group shell (tbody1) must not be reachable anywhere.
    reachable = _all_reachable_objgens(pdf)
    assert tbody1_key not in reachable
    tags = []
    doc = ptw._kids(pdf.Root.StructTreeRoot)[0]
    _walk_tags(doc, tags)
    assert tags.count("/TBody") == 1  # only the survivor's own TBody remains
    assert tags.count("/Table") == 1  # donor /Table gone, only survivor remains


# ── ragged-merge guard ──────────────────────────────────────────────────────

def test_differing_column_counts_refused_without_force():
    pdf = pikepdf.new()
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "a"), _cell(pdf, "b")])])  # 2 cols
    t1 = _table(pdf, [_row(pdf, [_cell(pdf, "x"), _cell(pdf, "y"), _cell(pdf, "z")])])  # 3 cols
    _doc_with_tables(pdf, [t0, t1])

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])
    assert result["tables_merged"] == 0
    assert result["rows_moved"] == 0
    assert result["donors_removed"] == 0
    assert len(result["skipped"]) == 1
    assert "column counts differ" in result["skipped"][0]["reason"]

    tables = ptw._find_tables(pdf)
    assert len(tables) == 2  # untouched


def test_differing_column_counts_merge_with_force():
    pdf = pikepdf.new()
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "a"), _cell(pdf, "b")])])  # 2 cols
    t1 = _table(pdf, [_row(pdf, [_cell(pdf, "x"), _cell(pdf, "y"), _cell(pdf, "z")])])  # 3 cols
    _doc_with_tables(pdf, [t0, t1])

    result = ptw.merge_tables(pdf, [{"indices": [0, 1], "force": True}])
    assert result["tables_merged"] == 1
    assert result["rows_moved"] == 1
    assert result["skipped"] == []

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    assert len(ptw._table_rows(tables[0])) == 2


# ── ascending-order guard ────────────────────────────────────────────────────
# Merge order IS reading order — a non-ascending merge silently scrambles a
# tagged table's reading order, and neither veraPDF (structural validity
# only) nor the fidelity gate (rendered pixels only, unchanged by a
# tag-tree-only merge) can catch it. See merge_tables' docstring.

def _three_frag_doc():
    pdf = pikepdf.new()
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "a1"), _cell(pdf, "a2")])])
    t1 = _table(pdf, [_row(pdf, [_cell(pdf, "b1"), _cell(pdf, "b2")])])
    t2 = _table(pdf, [_row(pdf, [_cell(pdf, "c1"), _cell(pdf, "c2")])])
    _doc_with_tables(pdf, [t0, t1, t2])
    return pdf


def test_non_ascending_order_refused_by_default():
    pdf = _three_frag_doc()

    result = ptw.merge_tables(pdf, [{"indices": [2, 0, 1]}])
    assert result["tables_merged"] == 0
    assert result["rows_moved"] == 0
    assert result["donors_removed"] == 0
    assert len(result["skipped"]) == 1
    reason = result["skipped"][0]["reason"]
    assert "not in ascending document order" in reason
    assert "reading order" in reason
    assert "veraPDF" in reason
    assert "fidelity gate" in reason

    tables = ptw._find_tables(pdf)
    assert len(tables) == 3  # untouched


def test_non_ascending_order_permitted_with_force():
    pdf = _three_frag_doc()

    result = ptw.merge_tables(pdf, [{"indices": [2, 0, 1], "force": True}])
    assert result["tables_merged"] == 1
    assert result["rows_moved"] == 2
    assert result["donors_removed"] == 2
    assert result["skipped"] == []

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    rows = ptw._table_rows(tables[0])
    texts = [ptw._cell_text(ptw._row_cells(r)[0]) for r in rows]
    assert texts == ["c1", "a1", "b1"]  # survivor (index 2) first, then given order


def test_ascending_order_unaffected_by_guard():
    pdf = _three_frag_doc()

    result = ptw.merge_tables(pdf, [{"indices": [0, 1, 2]}])
    assert result["tables_merged"] == 1
    assert result["skipped"] == []

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    rows = ptw._table_rows(tables[0])
    texts = [ptw._cell_text(ptw._row_cells(r)[0]) for r in rows]
    assert texts == ["a1", "b1", "c1"]


def test_two_index_descending_order_refused_by_default():
    """The two-fragment case: [1, 0] is non-ascending too, not just longer
    permutations."""
    pdf = _three_frag_doc()

    result = ptw.merge_tables(pdf, [{"indices": [1, 0]}])
    assert result["tables_merged"] == 0
    assert "not in ascending document order" in result["skipped"][0]["reason"]

    tables = ptw._find_tables(pdf)
    assert len(tables) == 3  # untouched


# ── refusal shapes: never raise ─────────────────────────────────────────────

def test_fewer_than_two_indices_refused():
    pdf = _two_frag_doc()
    result = ptw.merge_tables(pdf, [{"indices": [0]}])
    assert result["tables_merged"] == 0
    assert result["skipped"][0]["reason"] == "at least two indices required"
    assert len(ptw._find_tables(pdf)) == 2


def test_out_of_range_index_refused():
    pdf = _two_frag_doc()
    result = ptw.merge_tables(pdf, [{"indices": [0, 99]}])
    assert result["tables_merged"] == 0
    assert result["skipped"][0]["reason"] == "index out of range"
    assert len(ptw._find_tables(pdf)) == 2


def test_repeated_index_refused():
    pdf = _two_frag_doc()
    result = ptw.merge_tables(pdf, [{"indices": [0, 0]}])
    assert result["tables_merged"] == 0
    assert result["skipped"][0]["reason"] == "duplicate indices"
    assert len(ptw._find_tables(pdf)) == 2


def test_empty_specs_and_malformed_spec_shapes_never_raise():
    pdf = _two_frag_doc()
    result = ptw.merge_tables(pdf, [])
    assert result == {"tables_merged": 0, "rows_moved": 0, "donors_removed": 0, "skipped": []}

    result = ptw.merge_tables(pdf, [{"no_indices_key": True}, {"indices": "not a list"}])
    assert result["tables_merged"] == 0
    assert len(result["skipped"]) == 2
    assert len(ptw._find_tables(pdf)) == 2


# ── caption preservation ────────────────────────────────────────────────────

def test_donor_caption_survives_merge():
    pdf = pikepdf.new()
    survivor_caption = pdf.make_indirect(pikepdf.Dictionary(
        {"/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Caption"),
         "/K": pikepdf.Array([])}))
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "a"), _cell(pdf, "b")])], caption=survivor_caption)
    donor_caption = pdf.make_indirect(pikepdf.Dictionary(
        {"/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Caption"),
         "/K": pikepdf.Array([])}))
    t1 = _table(pdf, [_row(pdf, [_cell(pdf, "c"), _cell(pdf, "d")])], caption=donor_caption)
    _doc_with_tables(pdf, [t0, t1])
    donor_caption_key = donor_caption.objgen

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])
    assert result["tables_merged"] == 1

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    survivor = tables[0]
    captions = [k for k in ptw._kids(survivor) if str(k.get("/S", "")) == "/Caption"]
    assert len(captions) == 2  # both survivor's own and donor's caption present
    caption_keys = {c.objgen for c in captions}
    assert donor_caption_key in caption_keys  # donor's original caption object survives
    assert survivor_caption.objgen in caption_keys
    donor_caption_after = next(c for c in captions if c.objgen == donor_caption_key)
    assert donor_caption_after["/P"].objgen == survivor.objgen  # reparented onto survivor


def test_caption_only_donor_no_rows_still_moves_caption():
    """A donor with zero rows (just a caption) — merges cleanly, contributes
    0 rows, and its caption is not lost."""
    pdf = pikepdf.new()
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "a"), _cell(pdf, "b")])])
    donor_caption = pdf.make_indirect(pikepdf.Dictionary(
        {"/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Caption"),
         "/K": pikepdf.Array([])}))
    t1 = _table(pdf, [], caption=donor_caption)
    _doc_with_tables(pdf, [t0, t1])

    result = ptw.merge_tables(pdf, [{"indices": [0, 1], "force": True}])
    assert result["tables_merged"] == 1
    assert result["rows_moved"] == 0
    assert result["donors_removed"] == 1

    tables = ptw._find_tables(pdf)
    assert len(tables) == 1
    captions = [k for k in ptw._kids(tables[0]) if str(k.get("/S", "")) == "/Caption"]
    assert len(captions) == 1
    assert captions[0].objgen == donor_caption.objgen


# ── opt-in: tables absent from specs untouched ──────────────────────────────

def test_table_absent_from_spec_untouched():
    pdf = pikepdf.new()
    t0 = _table(pdf, [_row(pdf, [_cell(pdf, "a"), _cell(pdf, "b")])])
    t1 = _table(pdf, [_row(pdf, [_cell(pdf, "c"), _cell(pdf, "d")])])
    t2 = _table(pdf, [_row(pdf, [_cell(pdf, "e"), _cell(pdf, "f")])])
    _doc_with_tables(pdf, [t0, t1, t2])
    t2_key = t2.objgen

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])
    assert result["tables_merged"] == 1

    tables = ptw._find_tables(pdf)
    assert len(tables) == 2  # merged pair + untouched t2
    assert any(t.objgen == t2_key for t in tables)


# ── removed /Table truly unreachable ────────────────────────────────────────

def test_no_removed_table_reachable_after_merge():
    pdf = _two_frag_doc()
    tables_before = ptw._find_tables(pdf)
    donor_key = tables_before[1].objgen

    ptw.merge_tables(pdf, [{"indices": [0, 1]}])

    reachable = _all_reachable_objgens(pdf)
    assert donor_key not in reachable
    assert len(ptw._find_tables(pdf)) == 1


# ── ParentTree remap ─────────────────────────────────────────────────────

def _mcid_cell(pdf, mcid, pg):
    return pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/TD"),
        "/K": pikepdf.Array([mcid]), "/Pg": pg,
    }))


def _build_pdf_for_remap(path, with_parenttree=True):
    """1-page doc, 2 /Table fragments (2 cols each, matches for a clean
    merge): fragment 0 holds MCIDs 0/1, fragment 1 holds MCIDs 2/3. With
    with_parenttree=True, page.StructParents=0 and a 4-slot ParentTree
    /Nums array are wired, pre-pointing each MCID at its own cell (the
    realistic pre-merge state a tagger produces)."""
    pdf = pikepdf.new()
    page = pdf.add_blank_page(page_size=(200, 200))

    c0, c1 = _mcid_cell(pdf, 0, page.obj), _mcid_cell(pdf, 1, page.obj)
    row0 = _row(pdf, [c0, c1])
    t0 = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Table"),
        "/K": pikepdf.Array([row0]), "/Pg": page.obj,
    }))
    c2, c3 = _mcid_cell(pdf, 2, page.obj), _mcid_cell(pdf, 3, page.obj)
    row1 = _row(pdf, [c2, c3])
    t1 = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Table"),
        "/K": pikepdf.Array([row1]), "/Pg": page.obj,
    }))

    doc_elem = pdf.make_indirect(pikepdf.Dictionary({
        "/Type": pikepdf.Name("/StructElem"), "/S": pikepdf.Name("/Document"),
        "/K": pikepdf.Array([t0, t1]),
    }))
    t0["/P"] = doc_elem
    t1["/P"] = doc_elem

    root_dict = {
        "/Type": pikepdf.Name("/StructTreeRoot"),
        "/K": pikepdf.Array([doc_elem]),
    }
    if with_parenttree:
        nums_arr = pikepdf.Array([c0, c1, c2, c3])
        root_dict["/ParentTree"] = pikepdf.Dictionary({"/Nums": pikepdf.Array([0, nums_arr])})
        root_dict["/ParentTreeNextKey"] = 1
    pdf.Root.StructTreeRoot = pdf.make_indirect(pikepdf.Dictionary(root_dict))
    doc_elem["/P"] = pdf.Root.StructTreeRoot
    if with_parenttree:
        page.StructParents = 0
    pdf.save(path)


def test_parenttree_entries_still_point_at_correct_cell_after_merge(tmp_path):
    p = tmp_path / "remap.pdf"
    _build_pdf_for_remap(p)
    pdf = pikepdf.open(p)

    tables = ptw._find_tables(pdf)
    row1_cells = ptw._row_cells(ptw._table_rows(tables[1])[0])
    c2_key, c3_key = row1_cells[0].objgen, row1_cells[1].objgen

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])
    assert result["tables_merged"] == 1

    nums = pdf.Root.StructTreeRoot.ParentTree.Nums
    arr = nums[1]
    # moved MCIDs 2 and 3 still resolve to their own (reparented) cells
    assert arr[2].objgen == c2_key
    assert arr[3].objgen == c3_key


def test_no_structparents_merge_does_not_raise(tmp_path):
    p = tmp_path / "noparenttree.pdf"
    _build_pdf_for_remap(p, with_parenttree=False)
    pdf = pikepdf.open(p)

    result = ptw.merge_tables(pdf, [{"indices": [0, 1]}])
    assert result["tables_merged"] == 1
    assert len(ptw._find_tables(pdf)) == 1
