from aix.pipeline.assemble import build_index, segment

Q = [{"id": "q1", "label": "১", "alt_labels_json": [], "section_id": "A"},
     {"id": "q2", "label": "২", "alt_labels_json": [], "section_id": "B"},
     {"id": "q3", "label": "৩", "alt_labels_json": [], "section_id": "B"}]
U = {"q1": [{"id": "u1", "label": None}],
     "q2": [{"id": "u2a", "label": "ক"}, {"id": "u2b", "label": "খ"}],
     "q3": [{"id": "u3a", "label": "ক"}, {"id": "u3b", "label": "খ"}]}
QI, UI = build_index(Q, U)


def els(spec):
    """spec: list of (page, type, text, flags)."""
    out, sort = [], {}
    for i, (page, typ, text, *rest) in enumerate(spec):
        sort[page] = sort.get(page, -1) + 1
        out.append({"id": f"e{i}", "page_sort": page, "sort": sort[page], "type": typ, "text": text,
                    "q_label": text if typ == "heading" else None, "flags_json": rest[0] if rest else [],
                    "page_kind": "answer", "latex": None, "description": None, "table_json": None,
                    "candidates_json": None})
    return out


def by_unit(attempts):
    d = {}
    for a in attempts:
        d.setdefault(a.unit_id, []).append(a)
    return d


def test_out_of_order_and_continuation_across_pages():
    e = els([(1, "heading", "৩"), (1, "heading", "ক)"), (1, "line", "জ্ঞানমূলক উত্তর"),
             (1, "heading", "খ)"), (1, "line", "প্রথম লাইন"), (2, "line", "পরের পৃষ্ঠায় চলছে"),
             (2, "heading", "১"), (2, "line", "সংক্ষিপ্ত উত্তর")])
    d = by_unit(segment(e, Q, U, QI, UI))
    assert set(d) == {"u3a", "u3b", "u1"}
    assert len(d["u3b"][0].lines) == 3 and "[p2-l0]" in d["u3b"][0].lines[-1]
    assert d["u3a"][0].confidence >= 0.9


def test_line_label_and_sub_without_question_number():
    e = els([(1, "line", "উত্তর: ২ (ক) বাংলাদেশের রাজধানী ঢাকা"), (1, "line", "ঢাকা বড় শহর"),
             (1, "heading", "খ।"), (1, "line", "খ এর উত্তর")])
    d = by_unit(segment(e, Q, U, QI, UI))
    assert set(d) == {"u2a", "u2b"}
    assert d["u2a"][0].confidence == 0.85            # label found inside a line
    assert d["u2b"][0].confidence == 0.92            # resolved under the open question


def test_duplicate_attempt_and_struck_attempt():
    e = els([(1, "heading", "১"), (1, "line", "ভুল উত্তর", ["struck"]),
             (2, "heading", "১"), (2, "line", "সঠিক উত্তর")])
    atts = by_unit(segment(e, Q, U, QI, UI))["u1"]
    assert len(atts) == 2 and all("duplicate" in a.flags for a in atts)
    assert atts[0].struck == [False, True]           # heading not struck, content struck


def test_unknown_question_and_text_before_any_label_are_unmapped():
    e = els([(1, "line", "শুরুতে কিছু লেখা"), (1, "heading", "৯"), (1, "line", "এই উত্তরটি")])
    atts = segment(e, Q, U, QI, UI)
    assert [a.unit_id for a in atts] == [None, None]


def test_question_heading_without_parts_goes_to_first_part():
    e = els([(1, "heading", "২"), (1, "line", "ক এর উত্তর লেবেল ছাড়া"), (1, "heading", "খ)"), (1, "line", "খ উত্তর")])
    d = by_unit(segment(e, Q, U, QI, UI))
    assert "u2a" in d and d["u2a"][0].confidence == 0.70
    assert "u2b" in d


def test_continuation_keyword_appends_instead_of_duplicate():
    e = els([(1, "heading", "১"), (1, "line", "প্রথম অংশ"), (2, "heading", "২"), (2, "heading", "ক)"),
             (2, "line", "অন্য উত্তর"), (3, "line", "১ এর বাকি অংশ:"), ])
    # '১ এর বাকি অংশ' starts with a label-like '১' only if followed by marker; here it is ordinary text
    d = by_unit(segment(e, Q, U, QI, UI))
    assert len(d["u1"]) == 1


def test_rough_work_excluded():
    e = els([(1, "heading", "১"), (1, "line", "উত্তর"), (1, "line", "2+2=4", ["rough"])])
    a = segment(e, Q, U, QI, UI)[0]
    assert len(a.lines) == 2