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

"""MCP tools for workflow: status, outcome, brainstorming, direction."""

from __future__ import annotations

import logging
from uuid import UUID

from legal_mcp.services import db
from legal_mcp.services.lessons import (
    OUTCOME_LABELS_HE,
    VALID_OUTCOMES,
    canonical_outcome,
)
from legal_mcp.tools.envelope import empty, err, ok  # GAP-48: SSoT envelope

logger = logging.getLogger(__name__)


async def workflow_status(case_number: str) -> str:
    """סטטוס תהליך עבודה מלא לתיק - מסמכים, עיבוד, טיוטות.

    Args:
        case_number: מספר תיק הערר
    """
    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])
    docs = await db.list_documents(case_id)

    # Count chunks per document
    pool = await db.get_pool()
    async with pool.acquire() as conn:
        chunk_counts = await conn.fetch(
            "SELECT document_id, COUNT(*) as count FROM document_chunks WHERE case_id = $1 GROUP BY document_id",
            case_id,
        )
    chunk_map = {str(r["document_id"]): r["count"] for r in chunk_counts}

    doc_status = []
    for doc in docs:
        doc_status.append({
            "title": doc["title"],
            "type": doc["doc_type"],
            "extraction": doc["extraction_status"],
            "chunks": chunk_map.get(doc["id"], 0),
            "pages": doc.get("page_count"),
            # WS2 / #201: primary/secondary + whether claims were analysed.
            "is_primary": doc.get("is_primary", False),
            "claims_analyzed": doc.get("claims_analyzed", False),
        })

    # WS2 / #201: PRIMARY documents added after the analysis ran and not yet
    # included → the "not-analysed" condition that should prompt a merging
    # re-analysis (reanalyze_claims). Derived from is_primary + claims_extracted_at.
    pending_primary = await db.primary_docs_not_analyzed(case_id)
    primary_not_analyzed = [
        {"title": d["title"], "type": d["doc_type"]} for d in pending_primary
    ]

    # Check draft status
    from pathlib import Path
    from legal_mcp import config

    case_dir = config.find_case_dir(case_number)
    draft_path = case_dir / "drafts" / "decision.md"
    has_draft = draft_path.exists()
    draft_size = draft_path.stat().st_size if has_draft else 0

    status = {
        "case_number": case["case_number"],
        "title": case["title"],
        "status": case["status"],
        "documents": doc_status,
        "total_documents": len(docs),
        "total_chunks": sum(chunk_map.values()),
        "primary_docs_not_analyzed": primary_not_analyzed,
        "has_draft": has_draft,
        "draft_size_bytes": draft_size,
        "next_steps": _suggest_next_steps(case, docs, has_draft, primary_not_analyzed),
    }

    return ok(status)


def _suggest_next_steps(
    case: dict, docs: list, has_draft: bool,
    primary_not_analyzed: list | None = None,
) -> list[str]:
    """Suggest next steps based on case state."""
    steps = []
    doc_types = {d["doc_type"] for d in docs}

    if not docs:
        steps.append("העלה מסמכים לתיק (כתב ערר, תשובת ועדה)")
    else:
        if "appeal" not in doc_types:
            steps.append("העלה כתב ערר")
        if "response" not in doc_types:
            steps.append("העלה תשובת ועדה/משיבים")

    # WS2 / #201: a primary document was added but not yet analysed.
    if primary_not_analyzed:
        titles = ", ".join(d["title"] for d in primary_not_analyzed)
        steps.append(
            f"מסמך-עיקרי לא-נותח ({titles}) — הרץ ניתוח-מחדש מאחד (reanalyze_claims)"
        )

    pending = [d for d in docs if d["extraction_status"] == "pending"]
    if pending:
        steps.append(f"עיבוד {len(pending)} מסמכים ממתינים")

    if docs and not has_draft:
        steps.append("התחל ניסוח טיוטת החלטה (/draft-decision)")
    elif has_draft and case["status"] == "new":
        steps.append("סקור ועדכן את הטיוטה")
        steps.append("עדכן סטטוס ל-drafted")

    if case["status"] == "drafted":
        steps.append("סקירה סופית ועדכון סטטוס ל-reviewed")
    elif case["status"] == "reviewed":
        steps.append("אישור סופי ועדכון סטטוס ל-final")

    return steps


async def get_metrics(case_number: str = "") -> str:
    """מדדי הצלחה — KPIs לתיק ספציפי או דשבורד כולל.

    Args:
        case_number: מספר תיק (אם ריק — דשבורד כולל)
    """
    from legal_mcp.services import metrics

    if case_number:
        case = await db.get_case_by_number(case_number)
        if not case:
            return err(f"תיק {case_number} לא נמצא.")
        result = await metrics.get_case_metrics(UUID(case["id"]))
    else:
        result = await metrics.get_dashboard()

    return ok(result)


async def processing_status() -> str:
    """סטטוס כללי - מספר תיקים, מסמכים ממתינים לעיבוד."""
    pool = await db.get_pool()
    async with pool.acquire() as conn:
        case_count = await conn.fetchval("SELECT COUNT(*) FROM cases")
        doc_count = await conn.fetchval("SELECT COUNT(*) FROM documents")
        pending_count = await conn.fetchval(
            "SELECT COUNT(*) FROM documents WHERE extraction_status = 'pending'"
        )
        chunk_count = await conn.fetchval("SELECT COUNT(*) FROM document_chunks")
        corpus_count = await conn.fetchval("SELECT COUNT(*) FROM style_corpus")
        pattern_count = await conn.fetchval("SELECT COUNT(*) FROM style_patterns")

    return ok({
        "cases": case_count,
        "documents": doc_count,
        "pending_processing": pending_count,
        "chunks": chunk_count,
        "style_corpus_entries": corpus_count,
        "style_patterns": pattern_count,
    })


# ── Outcome & Brainstorming ───────────────────────────────────────

async def set_outcome(
    case_number: str,
    outcome: str,
    reasoning: str = "",
) -> str:
    """הזנת תוצאה לתיק ערר. יוצר רשומת החלטה ומפעיל סיעור מוחות אם אין נימוק.

    Args:
        case_number: מספר תיק הערר
        outcome: תוצאה — rejection (דחייה) / partial_acceptance (קבלה חלקית) /
            full_acceptance (קבלה מלאה). ערכי-legacy (rejected/accepted/partial) ממופים אוטומטית.
        reasoning: נימוק (אופציונלי). אם ריק — מפעיל סיעור מוחות.
    """
    from legal_mcp.services import brainstorm

    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    # GAP-51: accept legacy vocabulary (rejected/accepted/partial), store canonical.
    outcome = canonical_outcome(outcome)
    if outcome not in VALID_OUTCOMES:
        return err(f"תוצאה לא תקינה. אפשרויות: {', '.join(VALID_OUTCOMES)}")

    case_id = UUID(case["id"])

    # Create or update decision
    existing = await db.get_decision_by_case(case_id)
    if existing:
        await db.update_decision(
            UUID(existing["id"]),
            outcome=outcome,
            outcome_summary=reasoning[:200] if reasoning else "",
            outcome_reasoning=reasoning,
        )
        decision = await db.get_decision(UUID(existing["id"]))
    else:
        decision = await db.create_decision(
            case_id=case_id,
            outcome=outcome,
            outcome_summary=reasoning[:200] if reasoning else "",
            outcome_reasoning=reasoning,
        )

    # Update case status — aligned with the web /set-outcome endpoint:
    # reasoning given → direction is set; no reasoning → outcome only (brainstorm next).
    new_status = "direction_approved" if reasoning else "outcome_set"
    await db.update_case(case_id, status=new_status, expected_outcome=outcome)

    outcome_hebrew = OUTCOME_LABELS_HE.get(outcome, outcome)

    result = {
        "decision_id": decision["id"],
        "outcome": outcome,
        "outcome_hebrew": outcome_hebrew,
        "reasoning": reasoning,
        "has_reasoning": bool(reasoning),
    }

    if not reasoning:
        result["message"] = "לא סופק נימוק. הפעל /brainstorm כדי לקבל כיוונים אפשריים."
        result["next_step"] = "brainstorm"
    else:
        result["message"] = f"תוצאה נשמרה: {outcome_hebrew}. ניתן להתחיל כתיבת טיוטה."
        result["next_step"] = "draft"

    return ok(result)


async def brainstorm_directions(
    case_number: str,
) -> str:
    """סיעור מוחות — הצגת טענות מרכזיות והצעת 2-3 כיוונים אפשריים לנימוק.

    Args:
        case_number: מספר תיק הערר
    """
    from legal_mcp.services import brainstorm

    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])

    # Get existing decision for outcome
    decision = await db.get_decision_by_case(case_id)
    if not decision:
        return err("לא הוזנה תוצאה לתיק. הפעל set_outcome קודם.")

    outcome = decision.get("outcome", "")
    reasoning = decision.get("outcome_reasoning", "")

    directions = await brainstorm.generate_directions(case_id, outcome, reasoning)

    # Save brainstorm results to decision
    await db.update_decision(
        UUID(decision["id"]),
        direction_doc={"brainstorm": directions, "approved": False},
    )

    return ok(directions)


async def approve_direction(
    case_number: str,
    direction_index: int = 0,
    additional_notes: str = "",
) -> str:
    """אישור כיוון — יוצר מסמך כיוון מאושר. לא ניתן להתחיל כתיבת דיון בלי כיוון מאושר.

    Args:
        case_number: מספר תיק הערר
        direction_index: מספר הכיוון שנבחר (0 = ראשון)
        additional_notes: הערות נוספות
    """
    from legal_mcp.services import brainstorm

    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])
    decision = await db.get_decision_by_case(case_id)
    if not decision:
        return err("לא הוזנה תוצאה לתיק.")

    direction_data = decision.get("direction_doc") or {}
    brainstorm_result = direction_data.get("brainstorm", {})

    if not brainstorm_result.get("directions"):
        return err("לא בוצע סיעור מוחות. הפעל brainstorm_directions קודם.")

    direction_doc = brainstorm.build_direction_doc(
        outcome=decision.get("outcome", ""),
        reasoning=decision.get("outcome_reasoning", ""),
        directions_result=brainstorm_result,
        selected_direction=direction_index,
        additional_notes=additional_notes,
    )

    await db.update_decision(UUID(decision["id"]), direction_doc=direction_doc)

    return ok({"direction": direction_doc},
              message="כיוון אושר. ניתן להתחיל כתיבת טיוטה.")


async def ingest_final_version(
    case_number: str,
    file_path: str = "",
    final_text: str = "",
) -> str:
    """קליטת גרסה סופית (שדפנה חתמה). משווה לטיוטה ומחלצת לקחים.

    Args:
        case_number: מספר תיק הערר
        file_path: נתיב לקובץ הגרסה הסופית (PDF/DOCX)
        final_text: טקסט ישיר (אם אין קובץ)
    """
    from legal_mcp.services import learning_loop

    case = await db.get_case_by_number(case_number)
    if not case:
        return err(f"תיק {case_number} לא נמצא.")

    case_id = UUID(case["id"])

    # Extract text from file if provided
    if file_path and not final_text:
        from legal_mcp.services import extractor
        final_text, _, _ = await extractor.extract_text(file_path)

    if not final_text:
        return err("לא סופק טקסט — יש לספק file_path או final_text.")

    try:
        result = await learning_loop.process_final_version(case_id, final_text)
    except ValueError as e:
        return err(str(e))

    # Auto-ingest into internal committee decisions corpus (best-effort).
    # chair_name is resolved via the shared SoT (config.committee_chair_for_case)
    # — the SAME resolver the FastAPI upload path uses — so the two paths cannot
    # drift (INV-G2) and the DB chair constraint is never hit on an empty chair
    # (INV-G1: chair normalised at source). Failures are surfaced, not swallowed
    # (engineering rule §6 / feedback_silent_swallow): the result carries the
    # reason and final_learning_pipeline prints it.
    try:
        from legal_mcp import config
        from legal_mcp.services import internal_decisions as int_svc
        await int_svc.ingest_internal_decision(
            case_number=case_number,
            case_name=case.get("title", ""),
            decision_date=case.get("decision_date"),
            chair_name=config.committee_chair_for_case(case, case_number),
            district="ירושלים",
            practice_area=case.get("practice_area", ""),
            appeal_subtype=case.get("appeal_subtype", ""),
            text=final_text,
        )
        result["internal_corpus_ingested"] = True
    except Exception as e:
        logger.warning(
            "ingest_final_version: internal corpus ingestion failed (non-fatal): %s", e)
        result["internal_corpus_ingested"] = False
        result["internal_corpus_error"] = str(e)

    return ok(result)


# ── Chair feedback tools ──────────────────────────────────────────


async def record_chair_feedback(
    case_number: str,
    feedback_text: str,
    block_id: str = "block-yod",
    category: str = "missing_content",
    lesson_extracted: str = "",
) -> str:
    """תיעוד הערת יו"ר (דפנה) על טיוטת החלטה.

    Args:
        case_number: מספר תיק הערר
        feedback_text: ההערה של דפנה (מה חסר, מה לא נכון, מה צריך לשנות)
        block_id: הבלוק שההערה מתייחסת אליו (ברירת מחדל: block-yod)
        category: קטגוריה — missing_content/wrong_tone/wrong_structure/factual_error/style/other
        lesson_extracted: הלקח שהופק מההערה (אם ברור כבר)
    """
    case = await db.get_case_by_number(case_number)
    case_id = UUID(case["id"]) if case else None

    valid_categories = [
        "missing_content", "wrong_tone", "wrong_structure",
        "factual_error", "style", "other",
    ]
    if category not in valid_categories:
        return err(f"קטגוריה לא חוקית. אפשרויות: {', '.join(valid_categories)}")

    feedback_id = await db.record_chair_feedback(
        case_id=case_id,
        block_id=block_id,
        feedback_text=feedback_text,
        category=category,
        lesson_extracted=lesson_extracted,
    )

    # Auto-flow chair-authored STYLE feedback to the writer (closes the dead
    # chair_feedback→lesson chain — 27 feedback rows had produced 0 lessons). The
    # chair is the highest authority, so a style correction she writes flows
    # immediately — the chair IS the gate. It rides the SAME discussion_rules channel
    # promote uses (db.append_global_rule, G2), reaching every block, without the
    # style_corpus coupling decision_lessons require. SUBSTANCE feedback
    # (missing_content/factual_error/other) is case-specific → recorded only.
    # (INV-LRN1 graduated gate; 07-learning §1.2.)
    _STYLE_FB = {"style", "wrong_tone", "wrong_structure"}
    flowed = 0
    if lesson_extracted.strip() and category in _STYLE_FB:
        try:
            flowed = await db.append_global_rule(
                "discussion_rules", "universal", [lesson_extracted.strip()],
            )
        except Exception as e:
            logger.warning("chair-feedback auto-flow failed for %s: %s", case_number, e)

    msg = f"הערה נרשמה בהצלחה. קטגוריה: {category}."
    if flowed:
        msg += " הלקח (סגנון) זרם אוטומטית לכותב."
    return ok({
        "feedback_id": str(feedback_id),
        "flowed_to_writer": bool(flowed),
        "next_steps": [
            "כדי להפיק לקח מההערה, הפעל: analyze_chair_feedback",
            "כדי לסמן כמטופל: resolve_chair_feedback",
        ],
    }, message=msg)


_CURATOR_FINDING_CATEGORIES = {"style", "structure", "lexicon", "tabular", "general"}
# Agent tags ([סגנון]/[מבנה]/[לקסיקון משפטי]/[טבלאי]) → decision_lessons.category
_CURATOR_TAG_TO_CATEGORY = {
    "סגנון": "style", "מבנה": "structure",
    "לקסיקון משפטי": "lexicon", "לקסיקון": "lexicon", "טבלאי": "tabular",
}


async def record_curator_findings(case_number: str, findings: list[dict]) -> str:
    """לכידת ממצאי-האוצֵר כ-decision_lessons מובְנים (source='curator', proposed) — INV-LRN3.

    האוצֵר מזהה דפוסי-סגנון בקריאת הסופי; עד כה הם חיו רק כהערה ארעית בערוץ-התגובות
    (אובד, לא נסקר). כאן הם נתפסים מבנית כך שיופיעו בטאב ״אוצֵר״ ויעברו שער-יו"ר (INV-LRN1/G10).
    האוצֵר נשאר read-only על התוכן — הרישום הוא הצעה הממתינה לאישור, לא שינוי-קול.

    Args:
        case_number: מספר התיק הסופי (= decision_number בקורפוס-הסגנון).
        findings: רשימת ממצאים, כל אחד {"text": "...", "category"/"tag": "..."}.
                  category ∈ style/structure/lexicon/tabular/general (או tag עברי).
    """
    if not findings:
        return err("findings ריק — אין ממצאים לרשום.")
    corpus_id = await db.get_style_corpus_id_by_decision(case_number)
    if not corpus_id:
        return err(
            f"לא נמצאה רשומת style_corpus ל-{case_number} — ודא שהסופי נקלט לקורפוס-הסגנון "
            "(enroll_style_corpus) לפני רישום ממצאים."
        )
    # Dedup against lessons already on this corpus (any source) — re-running §A
    # must not pile duplicates (INV-LRN3 reliability).
    existing = {(_norm(r["lesson_text"])) for r in await db.list_decision_lessons(corpus_id)}
    written, skipped_dup, skipped_empty = [], 0, 0
    for f in findings:
        text = (f.get("text") or "").strip()
        if not text:
            skipped_empty += 1
            continue
        if _norm(text) in existing:
            skipped_dup += 1
            continue
        raw_cat = (f.get("category") or f.get("tag") or "general").strip()
        category = _CURATOR_TAG_TO_CATEGORY.get(raw_cat, raw_cat)
        if category not in _CURATOR_FINDING_CATEGORIES:
            category = "general"
        row = await db.add_decision_lesson(
            corpus_id,
            lesson_text=text,
            category=category,
            source="curator",
            created_by="curator",
            review_status="proposed",
        )
        if row:
            written.append(str(row["id"]))
            existing.add(_norm(text))
    return ok({
        "corpus_id": str(corpus_id),
        "written": len(written),
        "skipped_duplicate": skipped_dup,
        "skipped_empty": skipped_empty,
        "lesson_ids": written,
    }, message=(
        f"נרשמו {len(written)} ממצאי-אוצֵר (source=curator, ממתינים לשער-יו\"ר ב-/training). "
        f"{skipped_dup} כפילויות דולגו."
    ))


def _norm(s: str) -> str:
    """Normalize lesson text for dedup — collapse whitespace, strip."""
    return " ".join((s or "").split())


async def lesson_synthesize_pending(
    practice_area: str = "", category: str = "", apply: bool = True,
) -> str:
    """סינתזת-לקחים (#158 / INV-LRN8): ממזגת לקחי-סגנון חופפים ל"לקח-על" אחד עשיר ומוכלל,
    כך שהסט שזורם לכותב קטֵן ומשתבח (התקרה limit=15 מפסיקה לחתוך). מאשכלת לפי דמיון (cosine)
    בתוך shard של practice_area+category, ומסנתזת מעוגן-מקור (INV-AH) עם שער-drift.

    Args:
        practice_area: לצמצם ל-shard אחד (ריק = כל התחומים).
        category: style/structure/lexicon/tabular (ריק = כל הקטגוריות).
        apply: True = כותב (לקח-על approved + מקורות→superseded); False = dry-run.
    """
    from legal_mcp.services import lesson_synthesis
    shards = await lesson_synthesis.run_pending(practice_area, category, apply=apply)
    applied = sum(1 for s in shards for c in s["clusters"] if c.get("applied"))
    clusters = sum(len(s["clusters"]) for s in shards)
    return ok({
        "apply": apply,
        "shards": shards,
        "clusters_found": clusters,
        "synthesized": applied,
    }, message=(
        f"סינתזת-לקחים: {clusters} אשכולות ב-{len(shards)} shards · "
        f"{applied} לקחי-על {'נכתבו (approved)' if apply else 'דמו (dry-run)'}."
    ))


async def list_chair_feedback(
    case_number: str = "",
    category: str = "",
    unresolved_only: bool = True,
    limit: int = 100,
) -> str:
    """הצגת הערות יו"ר שתועדו, עם אפשרות סינון.

    Args:
        case_number: סינון לפי תיק (אם ריק — כל ההערות)
        category: סינון לפי קטגוריה
        unresolved_only: האם להציג רק הערות שלא טופלו (ברירת מחדל: כן)
        limit: תקרת תוצאות (INV-TOOL5 / GAP-53)
    """
    case_id = None
    if case_number:
        case = await db.get_case_by_number(case_number)
        if case:
            case_id = UUID(case["id"])

    feedbacks = await db.list_chair_feedback(
        case_id=case_id,
        category=category or None,
        unresolved_only=unresolved_only,
        limit=limit,
    )

    if not feedbacks:
        return empty("אין הערות שמתאימות לסינון.")

    items = []
    for fb in feedbacks:
        items.append({
            "id": str(fb["id"]),
            "case_id": str(fb["case_id"]) if fb["case_id"] else None,
            "block_id": fb["block_id"],
            "category": fb["category"],
            "feedback": fb["feedback_text"],
            "lesson": fb["lesson_extracted"],
            "resolved": fb["resolved"],
            "date": fb["created_at"].isoformat() if fb.get("created_at") else None,
        })

    return ok({
        "total": len(items),
        "feedbacks": items,
    })