backend/app/segmentation.py

"""Clause segmentation, hierarchy, defined terms and cross-references.

The brief calls this out as the real difficulty: "messy structure,
cross-references and defined terms". We segment on numbered headings,
build a parent/child hierarchy from numbering depth, harvest the
Definitions section, and resolve "the Supplier"-style references so
obligations name actual parties.
"""
from __future__ import annotations

import re

from .ingestion import IngestedDocument
from .schema import Clause, DefinedTerm

# "7." / "7.2" / "7.2.1" followed by a capitalized heading or sentence
CLAUSE_START = re.compile(r"^\s*(\d{1,2}(?:\.\d{1,2}){0,3})[.)]?\s+(?=[A-Z\"(])")
ARTICLE_START = re.compile(r"^\s*(ARTICLE|SECTION|Article|Section)\s+([IVXLC]+|\d+)\b[.:]?\s*")

# `"Confidential Information" means ...`
DEFINED_MEANS = re.compile(
    r"[\"“]([A-Z][\w\- /&]{1,60})[\"”]\s*\)?\s*"
    r"(?:means|shall mean|has the meaning|refers to)", re.UNICODE)
# `Globex Services Ltd ("Supplier")`
DEFINED_PAREN = re.compile(
    r"((?:[A-Z][\w.&,'-]*\s+){0,7}[A-Z][\w.&,'-]*)\s*\(\s*"
    r"(?:the\s+|each\s+a\s+|collectively,?\s+the\s+)?[\"“]([A-Z][\w\- /&]{1,40})[\"”]\s*\)",
    re.UNICODE)
CROSS_REF = re.compile(r"\b(?:Section|Clause|Article)s?\s+(\d{1,2}(?:\.\d{1,2}){0,3})")

# When a contract has no numbered/Article headings (common in real-world
# uploads), fall back to structure-agnostic segmentation so the classifier and
# risk engine still see clause-sized units instead of one giant block.
MAX_FALLBACK_CLAUSES = 80


def _blank_line_blocks(text: str) -> list[tuple[int, int]]:
    """Paragraph blocks separated by one or more blank lines."""
    blocks: list[tuple[int, int]] = []
    start = 0
    for m in re.finditer(r"\n[ \t]*\n+", text):
        if text[start:m.start()].strip():
            blocks.append((start, m.start()))
        start = m.end()
    if text[start:].strip():
        blocks.append((start, len(text)))
    return blocks


def _sentence_window_blocks(text: str, target: int = 500) -> list[tuple[int, int]]:
    """Group sentences into ~target-char windows — last resort for a wall of
    text with no paragraph breaks, so classification still has granularity."""
    bounds = [m.end() for m in re.finditer(r"[.;]\s+", text)]
    blocks: list[tuple[int, int]] = []
    s = 0
    for b in bounds:
        if b - s >= target:
            blocks.append((s, b))
            s = b
    if text[s:].strip():
        blocks.append((s, len(text)))
    return blocks or [(0, len(text))]


def _merge_tiny(text: str, blocks: list[tuple[int, int]], min_len: int = 80) -> list[tuple[int, int]]:
    """Fold very short fragments (addresses, signature lines) into the previous
    block so we don't spray the classifier with noise."""
    out: list[tuple[int, int]] = []
    for s, e in blocks:
        if out and len(text[s:e].strip()) < min_len:
            out[-1] = (out[-1][0], e)
        else:
            out.append((s, e))
    return out

PARTY_ROLES = {
    "client", "supplier", "vendor", "customer", "contractor", "provider",
    "licensor", "licensee", "company", "consultant", "buyer", "seller",
    "lessor", "lessee", "borrower", "lender", "service provider", "partner",
}
COMPANY_MARKERS = re.compile(
    r"\b(Inc|Ltd|LLC|L\.L\.C|Corp|Corporation|Company|GmbH|Limited|plc|S\.A|Pvt|LLP)\b")


def segment(doc: IngestedDocument) -> list[Clause]:
    """Split the document into clauses on numbered-heading boundaries."""
    lines: list[tuple[int, int, str]] = []  # (start, end, text)
    pos = 0
    for line in doc.text.split("\n"):
        lines.append((pos, pos + len(line), line))
        pos += len(line) + 1

    starts: list[tuple[int, str | None, str]] = []  # (line_idx, number, raw line)
    for i, (_s, _e, line) in enumerate(lines):
        m = CLAUSE_START.match(line)
        if m:
            starts.append((i, m.group(1), line))
            continue
        m = ARTICLE_START.match(line)
        if m:
            starts.append((i, m.group(2), line))

    clauses: list[Clause] = []

    def make_clause(cid: str, number: str | None, first_line: str,
                    start_char: int, end_char: int, level: int) -> Clause:
        raw = doc.text[start_char:end_char]
        lead = len(raw) - len(raw.lstrip())  # keep span aligned to visible text
        text = raw.strip()
        s = start_char + lead
        e = s + len(text) if text else end_char
        heading = _heading_from(first_line)
        return Clause(
            id=cid, number=number, heading=heading, text=text, level=level,
            span=doc.span(s, e),
            cross_references=sorted({m.group(1) for m in CROSS_REF.finditer(text)}),
        )

    if not starts:
        # No numbered/Article headings: segment structure-agnostically so the
        # rest of the pipeline still works on real-world / unstructured uploads.
        blocks = _blank_line_blocks(doc.text)
        if len(blocks) <= 1:
            blocks = _sentence_window_blocks(doc.text)
        blocks = _merge_tiny(doc.text, blocks)[:MAX_FALLBACK_CLAUSES]
        for i, (s, e) in enumerate(blocks):
            clauses.append(make_clause(f"c{i}", None, "", s, e, 1))
        if not clauses:
            clauses.append(make_clause("c0", None, "", 0, len(doc.text), 1))
        return clauses

    first_line_idx = starts[0][0]
    if lines[first_line_idx][0] > 0:
        pre_end = lines[first_line_idx][0]
        if doc.text[:pre_end].strip():
            c = make_clause("preamble", None, "", 0, pre_end, 1)
            c.heading = "Preamble"
            clauses.append(c)

    for n, (li, number, raw) in enumerate(starts):
        start_char = lines[li][0]
        end_char = lines[starts[n + 1][0]][0] if n + 1 < len(starts) else len(doc.text)
        level = (number.count(".") + 1) if number and number[0].isdigit() else 1
        clauses.append(make_clause(f"c{n + 1}", number, raw, start_char, end_char, level))

    # parent links from numbering depth
    stack: list[Clause] = []
    for c in clauses:
        while stack and stack[-1].level >= c.level:
            stack.pop()
        if stack:
            c.parent_id = stack[-1].id
        stack.append(c)
    return clauses


def _heading_from(first_line: str) -> str | None:
    body = CLAUSE_START.sub("", first_line, count=1)
    body = ARTICLE_START.sub("", body, count=1)
    body = body.strip()
    m = re.match(r"^([A-Z][^.:\n]{2,80})[.:]", body)
    if m:
        return m.group(1).strip()
    if body and body == body.upper() and len(body) < 80:
        return body.title()
    # heading on its own line: "5. Term and Renewal" — short title-case phrase
    if re.fullmatch(r"[A-Z][\w ,&/'-]{2,79}", body) and len(body.split()) <= 8:
        return body
    return None


def extract_defined_terms(doc: IngestedDocument, clauses: list[Clause]) -> list[DefinedTerm]:
    terms: dict[str, DefinedTerm] = {}

    def clause_at(off: int) -> str | None:
        for c in clauses:
            if c.span.start_char <= off < c.span.end_char:
                return c.id
        return None

    for m in DEFINED_PAREN.finditer(doc.text):
        entity, term = m.group(1).strip(), m.group(2).strip()
        if term not in terms:
            terms[term] = DefinedTerm(term=term, definition=entity,
                                      clause_id=clause_at(m.start()))
    for m in DEFINED_MEANS.finditer(doc.text):
        term = m.group(1).strip()
        tail = doc.text[m.end(): m.end() + 160].split(".")[0].strip()
        if term not in terms:
            terms[term] = DefinedTerm(term=term, definition=tail,
                                      clause_id=clause_at(m.start()))
    return list(terms.values())


def resolve_parties(terms: list[DefinedTerm]) -> dict[str, str]:
    """Map role terms ('Supplier') -> legal entity names ('Globex Services Ltd')."""
    parties: dict[str, str] = {}
    for t in terms:
        if t.term.lower() in PARTY_ROLES and COMPANY_MARKERS.search(t.definition):
            parties[t.term] = t.definition
    # fallback: role term defined via parenthetical against any proper noun
    for t in terms:
        if t.term.lower() in PARTY_ROLES and t.term not in parties and t.definition:
            parties[t.term] = t.definition
    return parties


def resolve_party_mention(sentence: str, parties: dict[str, str]) -> str | None:
    """Return the resolved entity for the party mentioned earliest in a sentence."""
    best: tuple[int, str] | None = None
    for role, entity in parties.items():
        m = re.search(rf"\b(?:the\s+)?{re.escape(role)}\b", sentence)
        if m and (best is None or m.start() < best[0]):
            best = (m.start(), f"{entity} ({role})")
    if best:
        return best[1]
    m = re.search(r"\b(Each [Pp]arty|Both [Pp]arties|Neither [Pp]arty)\b", sentence)
    if m:
        return m.group(1)
    return None