legal-ai/mcp-server/src/legal_mcp/tools/legal_arguments.py

"""MCP tools — aggregated legal arguments (claim de-duplication)."""

from __future__ import annotations

from uuid import UUID

from legal_mcp.services import argument_aggregator, claims_extractor, db
from legal_mcp.tools.envelope import empty, err, ok  # GAP-48: SSoT envelope

# Party labels reused across this module (display + impact diff).
_PARTY_HE = {
    "appellant": "עוררים",
    "respondent": "משיבים",
    "committee": "ועדה מקומית",
    "permit_applicant": "מבקשי היתר",
    "unknown": "צד לא מזוהה",
}


async def aggregate_claims_to_arguments(
    case_number: str,
    force: bool = False,
) -> str:
    """כינוס פרופוזיציות גולמיות לטיעונים משפטיים מובחנים.

    Args:
        case_number: מספר תיק הערר.
        force: True = למחוק טיעונים קיימים ולחשב מחדש.
    """
    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])
    result = await argument_aggregator.aggregate_claims_to_arguments(
        case_id, force=force,
    )
    result["case_number"] = case_number
    return ok(result)


async def get_legal_arguments(
    case_number: str,
    party: str = "",
) -> str:
    """שליפת טיעונים משפטיים מאוגדים לתיק.

    Args:
        case_number: מספר תיק הערר.
        party: סינון לפי צד (appellant/respondent/committee/permit_applicant).
               ריק = כל הצדדים.
    """
    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])
    args = await argument_aggregator.get_legal_arguments(case_id, party=party)

    if not args:
        return empty(
            "לא נמצאו טיעונים מאוגדים. הרץ aggregate_claims_to_arguments תחילה.",
            data={"case_number": case_number, "arguments": []},
        )

    # Group by party for nicer display.
    by_party: dict[str, list[dict]] = {}
    for a in args:
        label = _PARTY_HE.get(a["party"], a["party"])
        by_party.setdefault(label, []).append(a)

    return ok({
        "case_number": case_number,
        "total": len(args),
        "by_party": by_party,
    })


# ── Re-analysis with merge + impact diff (WS2 / #201) ───────────────

def _snapshot(args: list[dict]) -> dict:
    """Build a comparable snapshot of aggregated arguments for the impact diff.

    Captures, per party: the set of argument titles and a priority histogram —
    enough to show the chair what changed without dumping full bodies. The
    priority mix is the deterministic "recommendation" signal: it summarises the
    balance of threshold/substantive/procedural/relief arguments per side.
    """
    by_party: dict[str, dict] = {}
    for a in args:
        party = a.get("party", "unknown")
        bucket = by_party.setdefault(party, {"titles": [], "priorities": {}})
        bucket["titles"].append((a.get("argument_title") or "").strip())
        pr = a.get("priority", "substantive")
        bucket["priorities"][pr] = bucket["priorities"].get(pr, 0) + 1
    return {"total": len(args), "by_party": by_party}


def _impact_diff(before: dict, after: dict) -> dict:
    """Diff two argument snapshots → a chair-facing before↔after summary.

    Reports, per party, which argument titles were added/removed and how the
    priority mix (the "recommendation" balance) shifted. ``changed`` is a quick
    boolean the UI/chair can gate on. Nothing here is auto-applied — it is purely
    surfaced for the human gate (G10).
    """
    parties = sorted(set(before["by_party"]) | set(after["by_party"]))
    per_party: list[dict] = []
    any_change = False
    for party in parties:
        b = before["by_party"].get(party, {"titles": [], "priorities": {}})
        a = after["by_party"].get(party, {"titles": [], "priorities": {}})
        b_titles, a_titles = set(b["titles"]), set(a["titles"])
        added = sorted(a_titles - b_titles)
        removed = sorted(b_titles - a_titles)
        prio_before = dict(sorted(b["priorities"].items()))
        prio_after = dict(sorted(a["priorities"].items()))
        changed = bool(added or removed or prio_before != prio_after)
        any_change = any_change or changed
        per_party.append({
            "party": party,
            "party_he": _PARTY_HE.get(party, party),
            "count_before": len(b["titles"]),
            "count_after": len(a["titles"]),
            "added_arguments": added,
            "removed_arguments": removed,
            "priority_before": prio_before,
            "priority_after": prio_after,
            "changed": changed,
        })
    return {
        "changed": any_change,
        "total_before": before["total"],
        "total_after": after["total"],
        "by_party": per_party,
    }


async def reanalyze_claims(
    case_number: str,
    reanalyze_all_primary: bool = False,
) -> str:
    """ניתוח-מחדש מאחד של טענות — מחלץ ממסמכים-עיקריים חדשים/לא-נותחו בלבד,
    מאחד עם הטענות הקיימות (לא מוחק הכול), מריץ צבירה-מחדש, ומחזיר בדיקת-השפעה
    (diff טיעונים/המלצה לפני↔אחרי) ליו"ר.

    Args:
        case_number: מספר תיק הערר.
        reanalyze_all_primary: True = לחלץ מחדש מכל המסמכים-העיקריים (ולא רק
            מאלה שטרם-נותחו). ברירת-מחדל False = רק מסמכים-עיקריים חדשים/לא-נותחו.

    מנגנון האיחוד (לא force-delete): כל מסמך מחולץ דרך
    ``claims_extractor.extract_and_store_claims``, ש-``store_claims`` שלו מחליף רק
    את טענות *אותו* מסמך (לפי ``source_document``) — כך טענות ממסמכים שכבר-נותחו
    נשמרות. הצבירה-מחדש (``aggregate_claims_to_arguments(force=True)``) מחשבת את
    הטיעונים מחדש מתוך **מערך-הטענות המאוחד השלם** — מסלול-החישוב הקנוני, לא מסלול
    מקביל (G2).
    """
    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])

    # 1. BEFORE snapshot — current aggregated arguments (the impact baseline).
    before_args = await argument_aggregator.get_legal_arguments(case_id)
    before = _snapshot(before_args)

    # 2. Select PRIMARY documents to (re)extract. Default: only those not yet
    #    analysed (the not-analysed flag). reanalyze_all_primary widens to every
    #    primary doc. Source of truth for "primary" = is_primary (V47).
    if reanalyze_all_primary:
        all_docs = await db.list_documents(case_id)
        targets = [d for d in all_docs if d.get("is_primary")]
    else:
        targets = await db.primary_docs_not_analyzed(case_id)

    if not targets:
        return ok({
            "case_number": case_number,
            "status": "no_pending_documents",
            "message": "אין מסמכים-עיקריים חדשים/לא-נותחו. לא בוצע ניתוח-מחדש.",
            "documents_analyzed": [],
            "impact": _impact_diff(before, before),
        })

    # 3. Extract+merge claims from the selected documents only. Each call replaces
    #    only that document's claims (merge), never the whole case (no force-delete).
    analyzed: list[dict] = []
    for doc in targets:
        text = await db.get_document_text(UUID(doc["id"]))
        if not text:
            await db.mark_document_claims_extracted(UUID(doc["id"]), status="no_claims")
            analyzed.append({"document": doc["title"], "status": "empty_text", "total": 0})
            continue
        res = await claims_extractor.extract_and_store_claims(
            case_id=case_id,
            document_id=UUID(doc["id"]),
            text=text,
            doc_type=doc["doc_type"],
        )
        analyzed.append({"document": doc["title"], **res})

    # 4. Re-aggregate from the now-complete merged claim set. force=True is the
    #    canonical recompute (deletes only legal_arguments, NOT claims) — same
    #    path aggregate_claims_to_arguments always uses, not a fork.
    agg = await argument_aggregator.aggregate_claims_to_arguments(case_id, force=True)

    # 5. AFTER snapshot + impact diff for the chair.
    after_args = await argument_aggregator.get_legal_arguments(case_id)
    after = _snapshot(after_args)
    impact = _impact_diff(before, after)

    return ok({
        "case_number": case_number,
        "status": "completed",
        "documents_analyzed": analyzed,
        "aggregation": agg,
        "impact": impact,
    })