"""Build the excerpts Lev was trained on.

A decision is passed as a header line, its front matter, and its numbered paragraphs
({number: text}, with the text starting "[N] "). Each function returns the `state` string to pass to
`model.system_one(state=..., questions=...)`.

    header        "Smith v. Jones, 2025 ONCA 100 | Court of Appeal for Ontario (appeal court, Ontario) | decided 2025-02-03"
    front         the style of cause, coram, counsel and so on, as plain text
    paragraphs    {41: "[41] The respondent relies on ...", 42: "[42] ..."}
"""


def clip(text, limit):
    return text if len(text) <= limit else text[:limit].rsplit(" ", 1)[0] + " [...]"


def citation_pair(header, front, paragraphs, cited_name, cited_citation, citing_paragraphs, budget=3400):
    """Excerpt for the citator questions: treatment, history, reported, dissent.

    `citing_paragraphs` are the numbers of the paragraphs that name or cite the case.
    """
    order = sorted(paragraphs)
    index = {n: i for i, n in enumerate(order)}
    cited = [n for n in dict.fromkeys(citing_paragraphs) if n in paragraphs]
    if not cited:
        raise ValueError("none of the citing paragraphs is in the decision")
    chosen = list(cited)
    for n in cited:                                  # then one neighbour on each side
        for j in (index[n] - 1, index[n] + 1):
            if 0 <= j < len(order) and order[j] not in chosen:
                chosen.append(order[j])
    kept, used = [], 0
    for n in chosen:
        text = clip(paragraphs[n], 1300)
        if used + len(text) > budget and kept:
            continue
        kept.append(n)
        used += len(text)
    kept.sort(key=index.get)
    first = [n for n in order[:1] if n not in kept]
    last = [n for n in order[-2:] if n not in kept]
    lines = [f"Citing decision: {header}", f"Cited case: {cited_name}, {cited_citation}", "",
             "Front matter of the citing decision:", clip(front, 500), ""]
    if first:
        lines += ["Opening paragraph:", clip(paragraphs[first[0]], 400), ""]
    lines.append("Paragraphs that cite the case, with neighbours:")
    lines += [clip(paragraphs[n], 1300) for n in kept]
    if last:
        lines += ["", "Closing paragraphs:"] + [clip(paragraphs[n], 350) for n in last]
    return "\n".join(lines)


def decision(header, front, paragraphs):
    """Excerpt for the decision-level questions: outcome, decision_type, proceeding, area and the three yes/no ones."""
    order = sorted(paragraphs)
    first, last = order[:3], [n for n in order[-5:] if n not in order[:3]]
    lines = [f"Decision: {header}", "", "Front matter:", clip(front, 900), "", "Opening paragraphs:"]
    lines += [clip(paragraphs[n], 500) for n in first]
    if last:
        lines += ["", "Closing paragraphs:"] + [clip(paragraphs[n], 600) for n in last]
    return "\n".join(lines)


def passage(header, lead, paragraphs, numbers, budget=3600):
    """Excerpt for the passage-level questions.

    `lead` is ["Issue: ..."] for standard_of_review, or ["Legal test: ...", "Element: ..."] for element_finding.
    `numbers` are the paragraphs that deal with the issue or the element.
    """
    kept, used = [], 0
    for n in dict.fromkeys(numbers):
        if n not in paragraphs:
            continue
        text = clip(paragraphs[n], 1300)
        if used + len(text) > budget and kept:
            break
        kept.append(text)
        used += len(text)
    if not kept:
        raise ValueError("none of the paragraphs is in the decision")
    return "\n".join([f"Decision: {header}", *lead, "", "Paragraphs:", *kept])