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

"""MCP tools for decision drafting support."""

from __future__ import annotations

import json
from pathlib import Path
from uuid import UUID

from legal_mcp import config
from legal_mcp.services import db, embeddings, research_md
from legal_mcp.services.lessons import (
    CITATION_GUIDANCE,
    DECISION_TEMPLATES,
    DISCUSSION_RULES,
    GOLDEN_RATIOS,
    OPENING_STRATEGIES,
    PARAGRAPH_LENGTHS,
    SUMMARY_STRATEGIES,
    TRANSITION_PHRASES,
    VALID_OUTCOMES,
    format_ratios_comment,
    get_lessons_for_outcome,
)

# Fallback template for cases without expected_outcome
DECISION_TEMPLATE = """# החלטה

## בפני: דפנה תמיר, יו"ר ועדת הערר מחוז ירושלים

**ערר מספר:** {case_number}
**נושא:** {subject}
**העוררים:** {appellants}
**המשיבים:** {respondents}
**כתובת הנכס:** {property_address}

---

## א. רקע עובדתי

[תיאור הרקע העובדתי של הערר]

## ב. טענות העוררים

[סיכום טענות העוררים]

## ג. טענות המשיבים

[סיכום טענות המשיבים]

## ד. דיון והכרעה

[ניתוח משפטי]

## ה. מסקנה

[מסקנת הוועדה]

## ו. החלטה

[ההחלטה הסופית]

---
ניתנה היום, {date}
דפנה תמיר, יו"ר ועדת הערר
"""


async def get_style_guide() -> str:
    """שליפת דפוסי הסגנון של דפנה - נוסחאות, ביטויים אופייניים ומבנה, כולל לקחים מעשיים."""
    patterns = await db.get_style_patterns()

    result = "# מדריך סגנון - דפנה תמיר\n\n"

    # Part 1: DB-sourced patterns (from analyze_style)
    if patterns:
        type_names = {
            "opening_formula": "נוסחאות פתיחה",
            "transition": "ביטויי מעבר",
            "citation_style": "סגנון ציטוט",
            "analysis_structure": "מבנה ניתוח",
            "closing_formula": "נוסחאות סיום",
            "characteristic_phrase": "ביטויים אופייניים",
            "argument_flow": "זרימת טיעון",
            "evidence_handling": "התייחסות לראיות",
        }

        grouped: dict[str, list] = {}
        for p in patterns:
            pt = p["pattern_type"]
            if pt not in grouped:
                grouped[pt] = []
            grouped[pt].append({
                "text": p["pattern_text"],
                "context": p["context"],
                "frequency": p["frequency"],
            })

        for ptype, items in grouped.items():
            result += f"## {type_names.get(ptype, ptype)}\n\n"
            for item in items:
                result += f"- **{item['text']}** ({item['context']}, תדירות: {item['frequency']})\n"
            result += "\n"
    else:
        result += "_לא נמצאו דפוסים מקורפוס. יש להעלות החלטות ולהריץ /style-report._\n\n"

    # Part 2: Lessons-based guidance
    result += "---\n\n# לקחים מעשיים\n\n"

    # Universal discussion rules
    result += "## כללי דיון אוניברסליים\n\n"
    for rule in DISCUSSION_RULES["universal"]:
        result += f"- {rule}\n"
    result += "\n"

    # Citation technique
    result += "## טכניקת ציטוט\n\n"
    result += f"{CITATION_GUIDANCE}\n\n"

    # Transition phrases
    result += "## ביטויי מעבר (מהשוואת טיוטות)\n\n"
    for p in TRANSITION_PHRASES:
        ctx = p["context"]
        outcome_note = f" [בעיקר ב-{p['outcome']}]" if p["outcome"] else ""
        result += f"- **{p['phrase']}** — {ctx}{outcome_note}\n"
    result += "\n"

    # Paragraph lengths
    result += "## אורכי פסקאות מומלצים\n\n"
    result += "| סוג | מילים |\n|-----|------|\n"
    labels = {
        "claims": "טענות",
        "discussion_regular": "דיון רגיל",
        "discussion_with_citation": "דיון + ציטוט",
        "discussion_average": "ממוצע דיון",
    }
    for key, (lo, hi) in PARAGRAPH_LENGTHS.items():
        result += f"| {labels.get(key, key)} | {lo}-{hi} |\n"
    result += "\n"

    # Golden ratios
    result += "## יחסי הזהב (אחוזי סעיפים מסך ההחלטה)\n\n"
    result += "| סוג ערר | רקע | טענות | דיון | סיכום |\n"
    result += "|---------|------|-------|------|-------|\n"
    outcome_labels = {
        "rejection": "רישוי נדחה",
        "full_acceptance": "רישוי מתקבל",
        "partial_acceptance": "רישוי קבלה חלקית",
        "betterment_levy": "היטל השבחה",
    }
    for outcome in VALID_OUTCOMES:
        r = GOLDEN_RATIOS[outcome]
        result += (
            f"| {outcome_labels[outcome]} "
            f"| {r['background'][0]}-{r['background'][1]}% "
            f"| {r['claims'][0]}-{r['claims'][1]}% "
            f"| {r['discussion'][0]}-{r['discussion'][1]}% "
            f"| {r['summary'][0]}-{r['summary'][1]}% |\n"
        )
    result += "\n"

    # Opening and summary strategies
    result += "## אסטרטגיות פתיחה וסיכום לפי תוצאה\n\n"
    for outcome in VALID_OUTCOMES:
        opening = OPENING_STRATEGIES[outcome]
        summary = SUMMARY_STRATEGIES[outcome]
        result += f"### {outcome_labels[outcome]}\n"
        result += f"- **פתיחה:** {opening['description']} ({opening['paragraphs'][0]}-{opening['paragraphs'][1]} פסקאות)\n"
        result += f"- **סיכום ({summary['heading']}):** {summary['description']}\n\n"

    return result


async def draft_section(
    case_number: str,
    section: str,
    instructions: str = "",
) -> str:
    """הרכבת הקשר מלא לניסוח סעיף בהחלטה - כולל עובדות מהמסמכים, תקדימים רלוונטיים ודפוסי סגנון.

    Args:
        case_number: מספר תיק הערר
        section: סוג הסעיף (facts, appellant_claims, respondent_claims, legal_analysis, conclusion, ruling)
        instructions: הנחיות נוספות לניסוח
    """
    case = await db.get_case_by_number(case_number)
    if not case:
        return f"תיק {case_number} לא נמצא."

    case_id = UUID(case["id"])
    expected_outcome = case.get("expected_outcome", "")

    # 1. Get relevant chunks from case documents
    section_query = {
        "facts": "רקע עובדתי של התיק",
        "appellant_claims": "טענות העוררים",
        "respondent_claims": "טענות המשיבים",
        "legal_analysis": "ניתוח משפטי ודיון",
        "conclusion": "מסקנות",
        "ruling": "החלטה",
    }.get(section, section)

    query_emb = await embeddings.embed_query(section_query)
    case_chunks = await db.search_similar(
        query_embedding=query_emb, limit=10, case_id=case_id
    )

    # 2. Get similar sections from precedents
    precedent_chunks = await db.search_similar(
        query_embedding=query_emb, limit=5, section_type=section
    )
    # Filter out chunks from the same case
    precedent_chunks = [c for c in precedent_chunks if str(c["case_id"]) != case["id"]]

    # 3. Get style patterns
    style_patterns = await db.get_style_patterns()

    # Build context
    context = {
        "case": {
            "case_number": case["case_number"],
            "title": case["title"],
            "appellants": case["appellants"],
            "respondents": case["respondents"],
            "subject": case["subject"],
            "property_address": case["property_address"],
            "expected_outcome": expected_outcome,
        },
        "section": section,
        "instructions": instructions,
        "case_documents": [
            {
                "document": c["document_title"],
                "section_type": c["section_type"],
                "content": c["content"],
            }
            for c in case_chunks
        ],
        "precedents": [
            {
                "case_number": c["case_number"],
                "document": c["document_title"],
                "content": c["content"][:500],
            }
            for c in precedent_chunks[:3]
        ],
        "style_patterns": [
            {
                "type": p["pattern_type"],
                "text": p["pattern_text"],
            }
            for p in style_patterns[:15]
        ],
    }

    # 4. Add outcome-aware drafting guidance
    if expected_outcome and expected_outcome in VALID_OUTCOMES:
        lessons = get_lessons_for_outcome(expected_outcome)
        guidance: dict = {
            "outcome": expected_outcome,
            "golden_ratios": lessons["golden_ratios"],
            "citation_guidance": lessons["citation_guidance"],
        }

        # Section-specific guidance
        if section == "legal_analysis":
            guidance["discussion_rules"] = lessons["discussion_rules"]
            guidance["opening_strategy"] = lessons["opening_strategy"]
            guidance["transition_phrases"] = lessons["transition_phrases"]
            guidance["paragraph_lengths"] = lessons["paragraph_lengths"]
        elif section in ("conclusion", "ruling"):
            guidance["summary_strategy"] = lessons["summary_strategy"]
        elif section == "facts":
            guidance["target_ratio"] = lessons["golden_ratios"].get("background", "")
            guidance["paragraph_lengths"] = {"claims": lessons["paragraph_lengths"].get("claims", "")}
        elif section in ("appellant_claims", "respondent_claims"):
            guidance["target_ratio"] = lessons["golden_ratios"].get("claims", "")
            guidance["paragraph_lengths"] = {"claims": lessons["paragraph_lengths"].get("claims", "")}

        context["drafting_guidance"] = guidance

    return json.dumps(context, ensure_ascii=False, indent=2)


async def get_chair_directions(case_number: str) -> str:
    """שליפת עמדות יו"ר הוועדה (דפנה) על סוגיות הערר, לצורך יצירת direction_doc
    לכותב. קורא מ-analysis-and-research.md (שנוצר ע"י legal-analyst ומולא ע"י
    דפנה דרך ה-UI).

    מחזיר JSON עם סטטוס, כמה סוגיות מולאו וכמה עדיין ריקות, ורשימה של עמדות
    מובנות — ניתן להזריק ישירות כ-direction_doc לבלוק י (דיון) ולבלוק יא (סיכום).

    סטטוסים:
      missing  — הקובץ לא קיים
      empty    — הקובץ קיים אבל כל העמדות ריקות (טרם נקבעה דעה)
      partial  — חלק מהעמדות מולאו
      complete — כל העמדות מולאו

    אם המצב הוא `empty` או `missing` — הכותב צריך לעצור ולבקש מדפנה למלא
    את הקובץ דרך ה-UI לפני המשך הכתיבה.

    Args:
        case_number: מספר תיק הערר
    """
    case_dir = config.find_case_dir(case_number)
    file_path = case_dir / "documents" / "research" / "analysis-and-research.md"
    result = research_md.extract_chair_directions(file_path)
    return json.dumps(result, ensure_ascii=False, indent=2)


async def get_decision_template(case_number: str) -> str:
    """קבלת תבנית מבנית להחלטה מלאה עם פרטי התיק, מותאמת לסוג התוצאה הצפויה.

    Args:
        case_number: מספר תיק הערר
    """
    from datetime import date

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

    expected_outcome = case.get("expected_outcome", "")

    format_args = dict(
        case_number=case["case_number"],
        subject=case["subject"],
        appellants=", ".join(case.get("appellants", [])),
        respondents=", ".join(case.get("respondents", [])),
        property_address=case.get("property_address", ""),
        date=date.today().strftime("%d.%m.%Y"),
    )

    # Use outcome-specific template if available
    if expected_outcome and expected_outcome in DECISION_TEMPLATES:
        # Add ratio comments
        format_args["ratios_background"] = format_ratios_comment(expected_outcome, "background")
        format_args["ratios_claims"] = format_ratios_comment(expected_outcome, "claims")
        format_args["ratios_discussion"] = format_ratios_comment(expected_outcome, "discussion")
        format_args["ratios_summary"] = format_ratios_comment(expected_outcome, "summary")

        template = DECISION_TEMPLATES[expected_outcome].format(**format_args)

        # Add guidance header
        opening = OPENING_STRATEGIES[expected_outcome]
        summary = SUMMARY_STRATEGIES[expected_outcome]
        header = (
            f"<!-- תבנית מותאמת ל: {expected_outcome} -->\n"
            f"<!-- פתיחת דיון: {opening['description']} -->\n"
            f"<!-- סיכום: {summary['description']} -->\n\n"
        )
        return header + template
    else:
        # Fallback to generic template
        template = DECISION_TEMPLATE.format(**format_args)
        if not expected_outcome:
            template = (
                "<!-- לא הוגדרה תוצאה צפויה. הגדר expected_outcome בתיק לקבלת תבנית מותאמת. -->\n"
                f"<!-- ערכים אפשריים: {', '.join(VALID_OUTCOMES)} -->\n\n"
            ) + template
        return template


async def validate_decision(case_number: str) -> str:
    """בדיקת QA אוטומטית על ההחלטה — 6 בדיקות. אם נכשלת בדיקה קריטית — ייצוא חסום.

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

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

    case_id = UUID(case["id"])

    try:
        result = await qa_validator.validate_decision(case_id)
        return json.dumps(result, default=str, ensure_ascii=False, indent=2)
    except ValueError as e:
        return json.dumps({"status": "error", "message": str(e)}, ensure_ascii=False, indent=2)


async def export_docx(case_number: str, output_path: str = "") -> str:
    """ייצוא החלטה לקובץ DOCX מעוצב — גופן David, RTL, כותרות, מספור סעיפים.

    Args:
        case_number: מספר תיק הערר
        output_path: נתיב לשמירה (אופציונלי — ברירת מחדל: תיקיית התיק)
    """
    from legal_mcp.services import docx_exporter

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

    case_id = UUID(case["id"])

    try:
        path = await docx_exporter.export_decision(case_id, output_path or None)
        return json.dumps({
            "status": "completed",
            "path": path,
            "message": f"DOCX נוצר: {path}",
        }, ensure_ascii=False, indent=2)
    except ValueError as e:
        return json.dumps({
            "status": "error",
            "message": str(e),
        }, ensure_ascii=False, indent=2)


async def get_block_context(case_number: str, block_id: str, instructions: str = "") -> str:
    """קבלת הקשר מלא לכתיבת בלוק — ללא קריאה ל-API. Claude Code כותב את הבלוק.

    Args:
        case_number: מספר תיק הערר
        block_id: מזהה הבלוק (block-he, block-vav, ..., block-yod-bet)
        instructions: הנחיות נוספות
    """
    from legal_mcp.services import block_writer

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

    case_id = UUID(case["id"])
    try:
        ctx = await block_writer.get_block_context(case_id, block_id, instructions)
        return json.dumps(ctx, default=str, ensure_ascii=False, indent=2)
    except ValueError as e:
        return str(e)


async def save_block_content(case_number: str, block_id: str, content: str) -> str:
    """שמירת בלוק שנכתב ע"י Claude Code ב-DB.

    Args:
        case_number: מספר תיק הערר
        block_id: מזהה הבלוק
        content: הטקסט שנכתב
    """
    from legal_mcp.services import block_writer

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

    case_id = UUID(case["id"])
    try:
        result = await block_writer.save_block_content(case_id, block_id, content)
        return json.dumps(result, default=str, ensure_ascii=False, indent=2)
    except ValueError as e:
        return str(e)


async def analyze_style() -> str:
    """הרצת ניתוח סגנון על קורפוס ההחלטות של דפנה. מחלץ דפוסי כתיבה ושומר אותם."""
    from legal_mcp.services.style_analyzer import analyze_corpus

    result = await analyze_corpus()
    return json.dumps(result, ensure_ascii=False, indent=2)


async def write_block(
    case_number: str,
    block_id: str,
    instructions: str = "",
) -> str:
    """כתיבת בלוק יחיד בהחלטה. כותב ושומר ב-DB.

    Args:
        case_number: מספר תיק הערר
        block_id: מזהה הבלוק: block-alef, block-bet, block-gimel, block-dalet, block-he, block-vav, block-zayin, block-chet, block-tet, block-yod, block-yod-alef, block-yod-bet
        instructions: הנחיות נוספות לניסוח
    """
    from legal_mcp.services import block_writer

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

    case_id = UUID(case["id"])

    try:
        result = await block_writer.write_and_store_block(case_id, block_id, instructions)
        return json.dumps(result, default=str, ensure_ascii=False, indent=2)
    except ValueError as e:
        return str(e)


async def write_all_blocks(
    case_number: str,
    start_from: str = "block-alef",
    instructions: str = "",
) -> str:
    """כתיבת כל הבלוקים בהחלטה, בלוק אחרי בלוק. שומר כל בלוק מיד אחרי כתיבה.

    Args:
        case_number: מספר תיק הערר
        start_from: מאיזה בלוק להתחיל (ברירת מחדל: block-alef)
        instructions: הנחיות כלליות
    """
    from legal_mcp.services import block_writer
    from legal_mcp.services.block_writer import BLOCK_CONFIG

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

    case_id = UUID(case["id"])

    # Determine start index
    start_idx = BLOCK_CONFIG.get(start_from, {}).get("index", 1)

    results = []
    block_order = sorted(BLOCK_CONFIG.items(), key=lambda x: x[1]["index"])

    for bid, cfg in block_order:
        if cfg["index"] < start_idx:
            continue

        try:
            result = await block_writer.write_and_store_block(case_id, bid, instructions)
            results.append({
                "block_id": bid,
                "title": result["title"],
                "word_count": result["word_count"],
                "status": "completed",
            })
        except ValueError as e:
            results.append({
                "block_id": bid,
                "title": cfg["title"],
                "status": "error",
                "error": str(e),
            })
            # Stop on critical error (e.g., missing direction for block-yod)
            if "כיוון מאושר" in str(e):
                break

    total_words = sum(r.get("word_count", 0) for r in results)
    return json.dumps({
        "blocks": results,
        "total_words": total_words,
        "completed": sum(1 for r in results if r["status"] == "completed"),
    }, default=str, ensure_ascii=False, indent=2)