"""
Legal Writing Templates for Baymax
==================================
Structured templates that force you to fill in substance.
NOT AI-generated prose - these are scaffolds with placeholders.

Why templates beat drafting:
- Forces you to fill in substance (good for learning)
- Prevents garbage AI prose
- Guarantees proper structure
- Faster than staring at blank page
"""

from typing import List, Optional, Dict
from dataclasses import dataclass


# ============================================================================
# DATA STRUCTURES
# ============================================================================

@dataclass
class TemplateSection:
    """A section of a template with placeholder."""
    label: str
    placeholder: str
    hint: Optional[str] = None


# ============================================================================
# IRAC PARAGRAPH TEMPLATE
# ============================================================================

def generate_irac_template(
    thesis: Optional[str] = None,
    rule_name: Optional[str] = None,
    facts: Optional[List[str]] = None,
    precedent_case: Optional[str] = None,
    record_cites: Optional[List[str]] = None
) -> str:
    """
    Generate an IRAC paragraph template.

    Args:
        thesis: Your conclusion/thesis for this paragraph
        rule_name: Name of the rule being applied (e.g., "Tellabs test")
        facts: List of key facts to address
        precedent_case: Case for comparison/analogy
        record_cites: Record citations for facts

    Returns:
        Formatted template with placeholders
    """
    output = []
    output.append("=" * 70)
    output.append("IRAC PARAGRAPH TEMPLATE")
    output.append("=" * 70)

    # Topic Sentence
    output.append("\n[TOPIC SENTENCE - State your conclusion]")
    if thesis:
        output.append(f"  {thesis}.")
    else:
        output.append("  [YOUR THESIS: ________________________________]")

    output.append("")

    # Rule
    output.append("[RULE - Black letter law with citation]")
    if rule_name:
        output.append(f"  Under {rule_name}, [STATE THE RULE]. [CITATION].")
    else:
        output.append("  [STATE THE GOVERNING RULE]. [CITATION].")

    output.append("")

    # Rule Explanation
    output.append("[RULE EXPLANATION - What the rule means (1-2 sentences)]")
    output.append("  This standard requires [EXPLAIN WHAT RULE REQUIRES].")
    output.append("  Courts have interpreted this to mean [INTERPRETATION].")

    output.append("")

    # Application
    output.append("[APPLICATION - Apply rule to YOUR facts]")

    if facts:
        for i, fact in enumerate(facts, 1):
            cite = record_cites[i-1] if record_cites and i <= len(record_cites) else "[R. ___]"
            output.append(f"\n  Fact {i}: {fact}. {cite}.")
            output.append(f"  Analysis: [EXPLAIN WHY THIS FACT SATISFIES/VIOLATES THE RULE]")
    else:
        output.append("\n  Fact 1: [STATE FACT WITH RECORD CITE].")
        output.append("  Analysis: [EXPLAIN WHY THIS FACT MATTERS]")
        output.append("\n  Fact 2: [STATE FACT WITH RECORD CITE].")
        output.append("  Analysis: [EXPLAIN WHY THIS FACT MATTERS]")

    output.append("")

    # Comparison (optional)
    output.append("[COMPARISON - If relevant]")
    if precedent_case:
        output.append(f"  ANALOGY: Like in {precedent_case}, where [PRECEDENT FACT],")
        output.append("           here [YOUR CLIENT'S FACT].")
        output.append("  -- OR --")
        output.append(f"  DISTINCTION: Unlike {precedent_case}, where [PRECEDENT FACT],")
        output.append("               here [YOUR CLIENT'S FACT] because [REASON].")
    else:
        output.append("  ANALOGY: Like in [CASE], where [PRECEDENT FACT],")
        output.append("           here [YOUR CLIENT'S FACT].")
        output.append("  -- OR --")
        output.append("  DISTINCTION: Unlike [CASE], where [PRECEDENT FACT],")
        output.append("               here [YOUR CLIENT'S FACT] because [REASON].")

    output.append("")

    # Mini-Conclusion
    output.append("[MINI-CONCLUSION - Tie back to thesis]")
    if thesis:
        output.append(f"  Therefore, {thesis.lower()}.")
    else:
        output.append("  Therefore, [RESTATE YOUR THESIS].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# CREAC PARAGRAPH TEMPLATE (Appellate Style)
# ============================================================================

def generate_creac_template(
    conclusion: Optional[str] = None,
    rule_name: Optional[str] = None,
    facts: Optional[List[str]] = None
) -> str:
    """
    Generate a CREAC paragraph template (more formal appellate style).

    Args:
        conclusion: Your legal conclusion
        rule_name: Name of the governing standard
        facts: Key facts to address

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("CREAC PARAGRAPH TEMPLATE (Appellate Style)")
    output.append("=" * 70)

    # Conclusion
    output.append("\n[CONCLUSION - Direct statement]")
    if conclusion:
        output.append(f"  {conclusion}.")
    else:
        output.append("  [STATE YOUR LEGAL CONCLUSION DIRECTLY].")

    output.append("")

    # Rule
    output.append("[RULE - Governing legal standard]")
    if rule_name:
        output.append(f"  {rule_name} requires [STATE STANDARD]. [CITATION].")
    else:
        output.append("  [STATE THE GOVERNING LEGAL STANDARD]. [CITATION].")

    output.append("")

    # Explanation
    output.append("[EXPLANATION - Unpack the rule]")
    output.append("  This standard means [EXPLAIN MEANING].")
    output.append("  Courts apply this by examining [KEY FACTORS].")
    output.append("  The [CIRCUIT] has explained that [ELABORATION]. [CITE].")

    output.append("")

    # Application
    output.append("[APPLICATION - Apply to your facts]")
    if facts:
        for i, fact in enumerate(facts, 1):
            output.append(f"\n  {fact}. [R. ___].")
            output.append("  [ANALYSIS OF THIS FACT]")
    else:
        output.append("\n  [FACT 1 WITH RECORD CITE].")
        output.append("  [ANALYSIS]")
        output.append("\n  [FACT 2 WITH RECORD CITE].")
        output.append("  [ANALYSIS]")

    output.append("")

    # Conclusion (restate)
    output.append("[CONCLUSION - Restate with connector]")
    if conclusion:
        output.append(f"  Therefore, {conclusion.lower()}.")
    else:
        output.append("  Therefore, [RESTATE CONCLUSION].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# MULTI-ELEMENT TEST TEMPLATE
# ============================================================================

def generate_multielement_template(
    claim: Optional[str] = None,
    elements: Optional[List[str]] = None,
    party: str = "Plaintiff"
) -> str:
    """
    Generate a multi-element test template.

    Args:
        claim: The claim or cause of action
        elements: List of elements to prove
        party: Which party must prove (default: Plaintiff)

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("MULTI-ELEMENT TEST TEMPLATE")
    output.append("=" * 70)

    # Roadmap
    output.append("\n[ROADMAP]")
    if claim and elements:
        elem_list = ", ".join(f"({i+1}) {e}" for i, e in enumerate(elements))
        output.append(f"  To establish {claim}, {party} must prove: {elem_list}.")
        output.append("  [CITATION TO ELEMENTS].")
    else:
        output.append(f"  To establish [CLAIM], {party} must prove:")
        output.append("  (1) [ELEMENT 1], (2) [ELEMENT 2], (3) [ELEMENT 3].")
        output.append("  [CITATION].")

    output.append("")

    # Element analysis
    if elements:
        for i, element in enumerate(elements, 1):
            output.append("-" * 70)
            output.append(f"ELEMENT {i}: {element.upper()}")
            output.append("-" * 70)
            output.append(f"\n  [RULE for {element}]")
            output.append(f"  [STATE RULE FOR THIS ELEMENT]. [CITATION].")
            output.append("")
            output.append("  [YOUR FACTS]")
            output.append("  Here, [FACT 1]. [R. ___].")
            output.append("  [FACT 2]. [R. ___].")
            output.append("")
            output.append("  [ANALYSIS]")
            output.append(f"  These facts satisfy {element} because [EXPLAIN].")
            output.append("")
            output.append("  [MINI-CONCLUSION]")
            output.append(f"  Element {i} is satisfied.")
            output.append("")
    else:
        for i in range(1, 4):
            output.append("-" * 70)
            output.append(f"ELEMENT {i}: [ELEMENT NAME]")
            output.append("-" * 70)
            output.append(f"\n  [RULE for Element {i}]")
            output.append("  [STATE RULE]. [CITATION].")
            output.append("")
            output.append("  [YOUR FACTS]")
            output.append("  Here, [FACT]. [R. ___].")
            output.append("")
            output.append("  [ANALYSIS]")
            output.append("  [EXPLAIN WHY FACTS SATISFY ELEMENT].")
            output.append("")
            output.append("  [MINI-CONCLUSION]")
            output.append(f"  Element {i} is satisfied.")
            output.append("")

    # Overall Conclusion
    output.append("=" * 70)
    output.append("[OVERALL CONCLUSION]")
    if claim:
        output.append(f"  Because all elements are met, the Court should [GRANT/DENY] {party}'s")
        output.append(f"  [MOTION] regarding {claim}.")
    else:
        output.append(f"  Because all elements are met, the Court should [RELIEF REQUESTED].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# STATEMENT OF FACTS TEMPLATE
# ============================================================================

def generate_facts_template(
    parties: Optional[Dict[str, str]] = None,
    case_type: Optional[str] = None,
    key_dates: Optional[List[str]] = None
) -> str:
    """
    Generate a Statement of Facts template.

    Args:
        parties: Dict with "plaintiff" and "defendant" names
        case_type: Type of case (e.g., "securities fraud")
        key_dates: List of key dates to include

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("STATEMENT OF FACTS TEMPLATE")
    output.append("=" * 70)

    # Procedural Posture
    output.append("\n[PROCEDURAL POSTURE - 2-3 sentences]")
    if parties and case_type:
        output.append(f"  {parties.get('plaintiff', '[PLAINTIFF]')} brought this {case_type} action")
        output.append(f"  against {parties.get('defendant', '[DEFENDANT]')} on [DATE]. [R. ___].")
    else:
        output.append("  [PLAINTIFF] brought this [TYPE] action against [DEFENDANT]")
        output.append("  on [DATE]. [R. ___].")
    output.append("  The [district court/lower court] [RULING BELOW]. [R. ___].")
    output.append("  [CURRENT POSTURE - e.g., 'This appeal followed.']")

    output.append("")

    # Background
    output.append("[BACKGROUND - Scene setting, 3-5 sentences]")
    output.append("  [WHO ARE THE PARTIES]")
    if parties:
        output.append(f"  {parties.get('defendant', '[DEFENDANT]')} is a [DESCRIPTION]. [R. ___].")
        output.append(f"  {parties.get('plaintiff', '[PLAINTIFF]')} is a [DESCRIPTION]. [R. ___].")
    else:
        output.append("  [DEFENDANT] is a [DESCRIPTION]. [R. ___].")
        output.append("  [PLAINTIFF] is a [DESCRIPTION]. [R. ___].")
    output.append("")
    output.append("  [WHAT RELATIONSHIP/TRANSACTION IS AT ISSUE]")
    output.append("  [DESCRIBE RELATIONSHIP/BUSINESS CONTEXT]. [R. ___].")

    output.append("")

    # Key Facts (chronological)
    output.append("[KEY FACTS - Chronological, each with record cite]")
    if key_dates:
        for date in key_dates:
            output.append(f"\n  {date}: [WHAT HAPPENED]. [R. ___].")
    else:
        output.append("\n  [DATE 1]: [FACT]. [R. ___].")
        output.append("\n  [DATE 2]: [FACT]. [R. ___].")
        output.append("\n  [DATE 3]: [FACT]. [R. ___].")
        output.append("\n  [DATE 4]: [FACT]. [R. ___].")
        output.append("\n  [DATE 5]: [FACT]. [R. ___].")

    output.append("")

    # Harmful Consequences
    output.append("[HARMFUL CONSEQUENCES - if relevant]")
    output.append("  As a result, [PLAINTIFF] [SUFFERED WHAT HARM]. [R. ___].")
    output.append("  [QUANTIFY DAMAGES IF APPLICABLE]. [R. ___].")

    output.append("")

    # Bridge to Argument
    output.append("[BRIDGE TO ARGUMENT - 1 sentence]")
    output.append("  [TRANSITION SENTENCE CONNECTING FACTS TO YOUR LEGAL ARGUMENT]")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# CASE SYNTHESIS TEMPLATE (For Rule Sections)
# ============================================================================

def generate_synthesis_template(
    supreme_court_case: Optional[str] = None,
    circuit_cases: Optional[List[str]] = None,
    rule_topic: Optional[str] = None
) -> str:
    """
    Generate a case synthesis template for Rule sections.

    Args:
        supreme_court_case: Supreme Court case name if applicable
        circuit_cases: List of circuit court case names
        rule_topic: Topic/issue being synthesized

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("CASE SYNTHESIS TEMPLATE (Rule Section)")
    output.append("=" * 70)

    if rule_topic:
        output.append(f"\nTopic: {rule_topic}")
        output.append("")

    # Supreme Court (if applicable)
    output.append("[SUPREME COURT RULE - if applicable]")
    if supreme_court_case:
        output.append(f"  In {supreme_court_case}, the Supreme Court held that")
        output.append("  [STATE THE HOLDING]. [CITATION WITH PINPOINT].")
    else:
        output.append("  In [SUPREME COURT CASE], the Supreme Court held that")
        output.append("  [HOLDING]. [CITATION].")

    output.append("")

    # Circuit Court Elaboration
    output.append("[CIRCUIT COURT ELABORATION]")
    if circuit_cases and len(circuit_cases) > 0:
        output.append(f"  The [CIRCUIT] has explained that [ELABORATION].")
        output.append(f"  {circuit_cases[0]}, [CITATION].")
    else:
        output.append("  The [CIRCUIT] has explained that [ELABORATION/CLARIFICATION].")
        output.append("  [CASE], [CITATION].")

    output.append("")

    # Sub-rule or Exception
    output.append("[SUB-RULE OR EXCEPTION - if applicable]")
    if circuit_cases and len(circuit_cases) > 1:
        output.append(f"  Moreover, [ADDITIONAL RULE OR EXCEPTION].")
        output.append(f"  {circuit_cases[1]}, [CITATION].")
    else:
        output.append("  Moreover, [ADDITIONAL RULE OR EXCEPTION]. [CITATION].")

    output.append("")

    # Application Preview
    output.append("[APPLICATION PREVIEW - optional]")
    output.append("  Under this framework, [BRIEF PREVIEW OF HOW IT APPLIES TO YOUR CASE].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# STANDARD OF REVIEW TEMPLATE
# ============================================================================

def generate_standard_template(
    decision_type: Optional[str] = None,
    standard: Optional[str] = None
) -> str:
    """
    Generate a Standard of Review template.

    Args:
        decision_type: Type of decision being reviewed
        standard: The standard (de novo, clear error, abuse of discretion)

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("STANDARD OF REVIEW TEMPLATE")
    output.append("=" * 70)

    # Standard
    output.append("\n[STANDARD]")
    if decision_type and standard:
        output.append(f'  "This Court reviews {decision_type} {standard}."')
        output.append("  [CITATION].")
    else:
        output.append('  "This Court reviews [TYPE OF DECISION] [de novo / for clear error /')
        output.append('   for abuse of discretion]." [CITATION].')

    output.append("")

    # Explanation
    output.append("[EXPLANATION - What this means]")
    if standard:
        if "de novo" in standard.lower():
            output.append("  De novo review means the Court gives no deference to the lower court's")
            output.append("  legal conclusions and decides the issue independently. [CITATION].")
        elif "clear error" in standard.lower():
            output.append("  Clear error review means the Court will reverse only if left with")
            output.append("  a definite and firm conviction that a mistake was made. [CITATION].")
        elif "abuse" in standard.lower():
            output.append("  Abuse of discretion means the Court will reverse only if the lower")
            output.append("  court's decision was arbitrary, capricious, or clearly unreasonable. [CITATION].")
        else:
            output.append("  [EXPLAIN WHAT THIS STANDARD MEANS]. [CITATION].")
    else:
        output.append("  [EXPLAIN WHAT THIS STANDARD REQUIRES]. [CITATION].")

    output.append("")

    # Application
    output.append("[APPLICATION TO YOUR CASE]")
    output.append("  Here, [EXPLAIN WHY THE STANDARD FAVORS YOUR POSITION].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# CONCLUSION TEMPLATE
# ============================================================================

def generate_conclusion_template(
    relief: Optional[str] = None,
    reasons: Optional[List[str]] = None
) -> str:
    """
    Generate a Conclusion template.

    Args:
        relief: Specific relief requested
        reasons: List of key reasons (for numbered list)

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("CONCLUSION TEMPLATE")
    output.append("=" * 70)

    # Restate Standard
    output.append("\n[RESTATE STANDARD]")
    output.append("  Under [STANDARD OF REVIEW], this Court should [GRANT/DENY] because")
    output.append("  [ONE-SENTENCE SUMMARY].")

    output.append("")

    # Three Key Reasons
    output.append("[KEY REASONS - 3 points]")
    if reasons and len(reasons) >= 3:
        output.append(f"  First, {reasons[0]}.")
        output.append(f"  Second, {reasons[1]}.")
        output.append(f"  Third, {reasons[2]}.")
    else:
        output.append("  First, [REASON 1 IN ONE SENTENCE].")
        output.append("  Second, [REASON 2 IN ONE SENTENCE].")
        output.append("  Third, [REASON 3 IN ONE SENTENCE].")

    output.append("")

    # Request for Relief
    output.append("[REQUEST FOR RELIEF]")
    if relief:
        output.append(f"  The Court should therefore {relief}.")
    else:
        output.append("  The Court should therefore [SPECIFIC RELIEF REQUESTED].")

    output.append("")

    # Optional Policy Kicker
    output.append("[OPTIONAL POLICY KICKER - one sentence on broader implications]")
    output.append("  [POLICY REASON WHY THIS OUTCOME IS CORRECT].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# ROADMAP TEMPLATE
# ============================================================================

def generate_roadmap_template(
    num_parts: int = 3,
    relief: Optional[str] = None
) -> str:
    """
    Generate a brief roadmap/introduction template.

    Args:
        num_parts: Number of parts in the brief (default 3)
        relief: Relief requested

    Returns:
        Formatted template
    """
    output = []
    output.append("=" * 70)
    output.append("BRIEF ROADMAP TEMPLATE")
    output.append("=" * 70)

    output.append("\n[ROADMAP PARAGRAPH]")
    output.append("")

    parts = []
    for i in range(1, num_parts + 1):
        ordinal = {1: "First", 2: "Second", 3: "Third", 4: "Fourth", 5: "Fifth"}.get(i, f"Part {i}")
        parts.append(f"  {ordinal}, [PART {i} THESIS - one sentence].")

    output.append("  This brief proceeds in {} parts.".format(
        {2: "two", 3: "three", 4: "four", 5: "five"}.get(num_parts, str(num_parts))
    ))

    for part in parts:
        output.append(part)

    output.append("")
    if relief:
        output.append(f"  For these reasons, the Court should {relief}.")
    else:
        output.append("  For these reasons, the Court should [RELIEF REQUESTED].")

    output.append("\n" + "=" * 70)

    return "\n".join(output)


# ============================================================================
# MCP ENTRY POINTS
# ============================================================================

def generate_template_mcp(
    template_type: str,
    **kwargs
) -> str:
    """
    Main MCP entry point for template generation.

    Args:
        template_type: Type of template to generate
        **kwargs: Template-specific arguments

    Returns:
        Formatted template
    """
    templates = {
        "irac": generate_irac_template,
        "creac": generate_creac_template,
        "multielement": generate_multielement_template,
        "elements": generate_multielement_template,
        "facts": generate_facts_template,
        "statement_of_facts": generate_facts_template,
        "synthesis": generate_synthesis_template,
        "standard": generate_standard_template,
        "standard_of_review": generate_standard_template,
        "conclusion": generate_conclusion_template,
        "roadmap": generate_roadmap_template,
    }

    template_type = template_type.lower().replace(" ", "_").replace("-", "_")

    if template_type not in templates:
        available = ", ".join(sorted(templates.keys()))
        return f"Unknown template type: {template_type}\n\nAvailable templates: {available}"

    try:
        return templates[template_type](**kwargs)
    except TypeError as e:
        return f"Error generating template: {str(e)}\n\nCheck your arguments."


def list_templates() -> str:
    """List all available templates."""
    output = []
    output.append("AVAILABLE LEGAL WRITING TEMPLATES")
    output.append("=" * 70)
    output.append("")

    templates = [
        ("irac", "IRAC paragraph - Issue, Rule, Application, Conclusion"),
        ("creac", "CREAC paragraph - More formal appellate style"),
        ("multielement", "Multi-element test - For claims with multiple elements"),
        ("facts", "Statement of Facts - Structured fact statement"),
        ("synthesis", "Case synthesis - For Rule sections synthesizing multiple cases"),
        ("standard", "Standard of Review - De novo, clear error, abuse of discretion"),
        ("conclusion", "Conclusion - Summary with relief requested"),
        ("roadmap", "Brief roadmap - Introduction paragraph structure"),
    ]

    for name, desc in templates:
        output.append(f"  {name:15} - {desc}")

    output.append("")
    output.append("Use: generate_template(template_type='irac', thesis='...', facts=[...])")

    return "\n".join(output)


# ============================================================================
# TEST
# ============================================================================

if __name__ == "__main__":
    print(generate_irac_template(
        thesis="Plaintiff adequately pled scienter",
        rule_name="the Tellabs standard",
        facts=["CEO received weekly accuracy reports", "CFO sold stock before disclosure"],
        precedent_case="Higginbotham"
    ))

    print("\n\n")

    print(generate_multielement_template(
        claim="securities fraud under Section 10(b)",
        elements=["material misrepresentation", "scienter", "reliance", "damages"],
        party="Plaintiff"
    ))
