"""Resolve a table's footnote reference markers to their definition text.

A marker printed on a table cell points at a footnote, but the footnote's
text can be printed in three places, each needing its own resolution step:

1. Below the table, inside the correction agent's crop. The agent already
   transcribed those as marker-and-text pairs (`FootnoteDef`), so pairing is
   a normalized-key match — `(1)`, `1` and a superscript one are the same
   marker. The crop ends where the page does, so a footnote that continues on
   the next page reaches the agent cut mid-sentence; when the scan tier's line
   for the same marker opens with everything the agent read and continues
   past it, that longer line is taken as the definition instead.
2. Beyond the crop: later on the page or on a following page, including a
   continuation block (markers `a` through `j` printed overleaf). Those are
   found by scanning the document's text lines in reading order after the
   table. Acceptance requires coherence: a matching line on the table's own
   page is trusted (it sits right under the table), and so is a line inside
   a run of consecutive marker-opened lines (the shape of a real
   continuation block). A solitary line whose leading token happens to be a
   digit, found on a distant page, proves nothing and is refused.
3. Nowhere as a footnote at all: a cross-reference to a NAMED SECTION of
   the document — `(Note 12)`, `Schedule II`, `(Addendum 3)` — points at a
   whole section, not a footnote line. WHICH markers are section references
   is the correction agent's judgement, made from the page image and carried
   on each mark's `kind`; nothing here infers it from the marker's wording.
   A section reference resolves to a POINTER — the heading that opens with
   the reference as written, plus its page — never to inlined text.

Everything here is pure and deterministic. The demand-driven lookup agent
(`quber.agents.footnote_lookup`) is the third tier for markers this module
leaves unresolved; callers invoke it separately so the common case costs no
model call. Every marker either resolves or lands in `unresolved`, and every
marked definition either matches a marker or lands in `unreferenced` —
nothing is silently dropped. Unmarked general notes are not exceptions: they
carry no marker to match, and the consumer attaches them at table level.
"""

from __future__ import annotations

import re
from typing import Collection, List, Literal, Optional, Sequence, Tuple

from pydantic import BaseModel, Field

from quber.agents.llm_client import FootnoteDef


class FootnoteResolution(BaseModel):
    """The complete resolution for one table's markers, with its exceptions."""

    resolved: List["ResolvedFootnote"] = Field(default_factory=list)
    unresolved: List[str] = Field(
        default_factory=list,
        description="Markers (as printed) whose definition was not found in any tier",
    )
    unreferenced: List[FootnoteDef] = Field(
        default_factory=list,
        description="Marked definitions that no marker on the table points at",
    )

    def text_for(self, marker: str) -> Optional[str]:
        """The resolved text for a printed marker, matched by canonical key."""
        key = canonical_marker(marker)
        for r in self.resolved:
            if canonical_marker(r.marker) == key:
                return r.text
        return None


class ResolvedFootnote(BaseModel):
    """One marker paired with the text a reader should see for it."""

    marker: str = Field(description="The marker as printed on the table")
    text: str = Field(description="The definition text, or the pointer line for a Notes reference")
    source: Literal["table", "scan", "note_pointer", "lookup", "sibling"] = Field(
        description=(
            "table: paired in the agent's own crop. scan: found by the reading-order "
            "scan beyond the crop. note_pointer: a Notes-section reference resolved to "
            "its heading and page. lookup: found by the demand-driven lookup agent. "
            "sibling: a definition printed under a neighboring table of the same series."
        )
    )


# Superscript glyphs normalize to their plain forms so a superscript one and a
# printed '(1)' key the same.
_SUPERSCRIPTS = str.maketrans("⁰¹²³⁴⁵⁶⁷⁸⁹ᵃᵇᶜᵈᵉᶠᵍʰⁱʲᵏ", "0123456789abcdefghijk")


# The leading token of a printed footnote line. Digits and symbols may stand
# bare ('2 Includes...'); letters must be set off by parentheses or trailing
# punctuation ('a.' / '(a)'), because an ordinary prose line also starts with
# a short word and would otherwise read as a marker.
_LEAD_MARKER_RE = re.compile(
    r"^(?:"
    r"\((?P<paren>\d{1,2}|[A-Za-z]{1,2}|[*†‡§]{1,3})\)"
    r"|(?P<bare>\d{1,2}|[*†‡§]{1,3})"
    r"|(?P<punct>[A-Za-z]{1,2})[.):]"
    r")\s+"
)


def canonical_marker(marker: str) -> str:
    """A marker's identity independent of rendering: '(1)', '1' and a
    superscript one all key to '1'; letters casefold."""
    s = re.sub(r"\s+", "", marker or "").translate(_SUPERSCRIPTS)
    return s.strip("()[]").rstrip(".:").casefold()


def split_leading_marker(line: str) -> Optional[Tuple[str, str]]:
    """(canonical marker, rest of line) when `line` opens the way a printed
    footnote does; None otherwise. The rest must be non-empty — a marker
    with nothing after it defines nothing."""
    m = _LEAD_MARKER_RE.match(line.strip())
    if not m:
        return None
    token = m.group("paren") or m.group("bare") or m.group("punct")
    rest = line.strip()[m.end() :].strip()
    if not rest:
        return None
    return canonical_marker(token), rest


def resolve_footnotes(
    markers: Sequence[str],
    footnotes: Sequence[FootnoteDef],
    table_page: int,
    trailing_lines: Sequence[Tuple[str, int]] = (),
    headings: Sequence[Tuple[str, int]] = (),
    section_keys: Collection[str] = frozenset(),
) -> FootnoteResolution:
    """Resolve every distinct marker on one table, in tier order.

    `markers` are the table's printed reference markers (the located marks
    plus the catalogued refs). `footnotes` are the agent's in-crop pairs.
    `trailing_lines` are the document's text lines in reading order after the
    table, each with its 1-indexed page — the scan tier's universe.
    `headings` are the document's section headings with their pages.
    `section_keys` are the canonical keys of the markers the correction agent
    judged to be cross-references to a named section of the document — that
    judgement is the agent's, made from the page image; nothing here infers
    it from the marker's wording or shape. A section reference resolves to a
    heading pointer or stays unresolved; it is never hunted as a footnote.
    Returns the resolution plus both exception lists; markers still
    unresolved here are the lookup agent's demand.
    """
    defs_by_key: dict[str, FootnoteDef] = {}
    for d in footnotes:
        key = canonical_marker(d.marker)
        if key:
            defs_by_key.setdefault(key, d)

    parsed_lines = [(split_leading_marker(text), page) for text, page in trailing_lines]
    # Runs of consecutive marker-opened lines: the shape of a continuation
    # block. block_len[i] is the length of the run line i belongs to.
    block_len = [0] * len(parsed_lines)
    i = 0
    while i < len(parsed_lines):
        if parsed_lines[i][0] is None:
            i += 1
            continue
        j = i
        while j < len(parsed_lines) and parsed_lines[j][0] is not None:
            j += 1
        for k in range(i, j):
            block_len[k] = j - i
        i = j

    resolved: List[ResolvedFootnote] = []
    unresolved: List[str] = []
    seen_keys: set[str] = set()
    for printed in markers:
        key = canonical_marker(printed)
        if not key or key in seen_keys:
            continue
        seen_keys.add(key)

        if key in section_keys:
            pointer = _section_pointer(printed, headings)
            if pointer:
                resolved.append(ResolvedFootnote(marker=printed, text=pointer, source="note_pointer"))
            else:
                unresolved.append(printed)
            continue

        candidates = _scan_candidates(key, parsed_lines, block_len, table_page)

        if key in defs_by_key:
            in_crop = defs_by_key[key].text
            # The crop ends where the page does. A footnote whose text runs on
            # to the next page reaches the agent cut mid-sentence, while the
            # parse stitches both pages into one line. When a line the scan
            # accepts for the same marker opens with everything the agent read
            # and continues past it, that line is the complete definition. The
            # agent's own text also rides in the scan universe (the graft
            # attaches it to the table), so the first match is often the same
            # cut text; every accepted line is tried.
            longer = next((c for c in candidates if _extends(in_crop, c)), None)
            if longer is not None:
                resolved.append(ResolvedFootnote(marker=printed, text=longer, source="scan"))
            else:
                resolved.append(ResolvedFootnote(marker=printed, text=in_crop, source="table"))
            continue

        if candidates:
            resolved.append(ResolvedFootnote(marker=printed, text=candidates[0], source="scan"))
            continue

        # A marker the agent judged a footnote but that no tier defines can
        # still be a mislabelled section reference. The test is the same one
        # section references pass: a document heading literally opens with
        # the printed marker. Restricted to word-plus-number markers
        # ('Note 18') — a bare '(1)' or '(a)' can never fall through to a
        # heading, so a genuine footnote marker with a missing definition
        # stays unresolved and flagged.
        if (
            any(ch.isalpha() for ch in printed)
            and any(ch.isdigit() for ch in printed)
            and (pointer := _section_pointer(printed, headings))
        ):
            resolved.append(ResolvedFootnote(marker=printed, text=pointer, source="note_pointer"))
        else:
            unresolved.append(printed)

    unreferenced = [
        d for d in footnotes if canonical_marker(d.marker) and canonical_marker(d.marker) not in seen_keys
    ]
    return FootnoteResolution(resolved=resolved, unresolved=unresolved, unreferenced=unreferenced)


def absorb_lookup(
    resolution: FootnoteResolution,
    found: Sequence[FootnoteDef],
    source: Literal["lookup", "sibling"] = "lookup",
) -> FootnoteResolution:
    """Fold externally found definitions into a resolution.

    Each found pair matching a still-unresolved marker resolves it under
    `source`; markers not found stay unresolved, and the rest of the
    resolution is unchanged."""
    by_key = {canonical_marker(d.marker): d for d in found if canonical_marker(d.marker)}
    resolved = list(resolution.resolved)
    unresolved: List[str] = []
    for printed in resolution.unresolved:
        d = by_key.get(canonical_marker(printed))
        if d is not None:
            resolved.append(ResolvedFootnote(marker=printed, text=d.text, source=source))
        else:
            unresolved.append(printed)
    return FootnoteResolution(
        resolved=resolved, unresolved=unresolved, unreferenced=list(resolution.unreferenced)
    )


def _extends(shorter: str, longer: str) -> bool:
    """Whether `longer` opens with the whole of `shorter` and continues past
    it, compared word by word with case, spacing and punctuation dropped so
    the two extractors' renderings of the same printed line agree. A hyphen
    dropped at a line break ('interest-only' against 'interestonly') is the
    remaining difference, so the comparison also joins the words."""
    a = _words(shorter)
    b = _words(longer)
    if not a or len(b) <= len(a):
        return False
    if b[: len(a)] == a:
        return True
    joined_a = "".join(a)
    joined_b = "".join(b)
    return len(joined_b) > len(joined_a) and joined_b.startswith(joined_a)


def _words(text: str) -> List[str]:
    return re.findall(r"[a-z0-9]+", text.casefold())


def _scan_candidates(
    key: str,
    parsed_lines: Sequence[Tuple[Optional[Tuple[str, str]], int]],
    block_len: Sequence[int],
    table_page: int,
) -> List[str]:
    """The scan tier: every coherent line opening with `key`, in reading order.

    Coherent means the line is on the table's own page (it sits right under
    the table), or it belongs to a run of two or more consecutive
    marker-opened lines (a continuation block), or its marker is not purely
    numeric (a bare digit is the one leading token common prose also
    produces). The remaining case — a solitary numeric match on a distant
    page — is refused. The first entry is the scan tier's own answer; the rest
    matter only when the agent's in-crop text is being checked for a longer
    rendering of the same footnote.
    """
    found: List[str] = []
    for (parsed, page), run in zip(parsed_lines, block_len, strict=True):
        if parsed is None or parsed[0] != key:
            continue
        if page == table_page or run >= 2 or not key.isdigit():
            found.append(parsed[1])
    return found


def _section_pointer(printed: str, headings: Sequence[Tuple[str, int]]) -> Optional[str]:
    """The pointer line for a section reference: the target's heading and page.

    The reference as the agent wrote it, normalized (spacing, punctuation,
    case dropped), must be the opening of a heading — 'Note 16' opens
    'NOTE 16. Commitments and Contingencies', 'Schedule II' opens
    'Schedule II — Valuation and Qualifying Accounts'. Two printed forms of
    the same target also agree: a reference written with a word where the
    heading opens with the bare number (or the reverse) matches on the
    number — 'Note 10' points at '10. Contingencies' — but only when the
    reference itself is a word-plus-number, never for a bare digit, and only
    at a digit boundary so '10' cannot open '104'. A reference naming
    several targets ('Notes 1 and 4') resolves when every number it names
    finds a heading; the pointer then lists them all. None when no heading
    matches — the reference then stays unresolved and flagged rather than
    pointing at nothing."""
    key = _heading_key(printed)
    if not key:
        return None
    for text, page in headings:
        if _heading_key(text).startswith(key):
            return f"See {text.strip()} (page {page})"

    lead = re.match(r"([a-z]+)\d", key)
    if not lead:
        return None
    ref_word = lead.group(1)
    pointers = []
    for number in re.findall(r"\d+", key):
        hit = None
        for text, page in headings:
            m = re.match(r"([a-z]*)(\d+)", _heading_key(text))
            if m is None or m.group(2) != number:
                continue
            hword = m.group(1)
            if hword and not (hword.startswith(ref_word) or ref_word.startswith(hword)):
                continue
            hit = f"See {text.strip()} (page {page})"
            break
        if hit is None:
            return None
        pointers.append(hit)
    return "; ".join(pointers) if pointers else None


def _heading_key(text: str) -> str:
    """Text reduced to its letters and digits, casefolded — the form in which
    a printed reference and its target heading agree regardless of spacing or
    punctuation."""
    return "".join(ch for ch in text if ch.isalnum()).casefold()