"""
Advanced Legal Analysis Tools
=============================
Tier 1 features for making this tool a WEAPON for brief writing:
1. Rule Synthesis - Extract and combine black letter rules
2. Argument Extraction - Identify arguments and supporting authority
3. Fact Pattern Analysis - Extract key facts and procedural posture
4. Case Comparison - Side-by-side analysis with distinction language
5. Memo Draft Generator - Auto-populate IRAC structure
"""

import re
from typing import Dict, List, Tuple, Optional, Set
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict


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

@dataclass
class LegalRule:
    """Represents a black letter rule extracted from a case"""
    text: str
    source_case: str
    rule_type: str  # "holding", "test", "elements", "factors", "standard"
    elements: List[str] = field(default_factory=list)
    context: str = ""


@dataclass
class Argument:
    """Represents a discrete legal argument"""
    claim: str
    supporting_authority: List[str]
    reasoning: str
    strength: int  # 1-5 based on authority depth
    argument_type: str  # "main", "alternative", "counter"
    heading: str = ""  # Section heading if from brief structure


@dataclass
class FactPattern:
    """Represents extracted facts from a case or brief"""
    procedural_posture: str
    key_facts: List[str]
    background_facts: List[str]
    legally_significant_facts: List[str]
    parties: Dict[str, str]  # role -> name


@dataclass
class CaseComparison:
    """Represents comparison between two cases or case vs. facts"""
    similarities: List[str]
    differences: List[str]
    distinguishing_factors: List[str]
    suggested_language: List[str]


@dataclass
class MemoDraft:
    """Represents a draft legal memo structure"""
    question_presented: str
    brief_answer: str
    facts_template: str
    discussion_irac: Dict[str, str]  # section -> content
    conclusion: str


# ============================================================================
# RULE EXTRACTION PATTERNS
# ============================================================================

# Patterns that indicate a legal rule/holding
RULE_PATTERNS = [
    # ===== HOLDINGS (explicit markers) =====
    # Supreme Court "Held:" format
    (r'[Hh]eld[:\s]+(?:that\s+)?(.{50,500}?)(?:\.|$)', 'holding'),
    # Explicit "we hold" variations
    (r'[Ww]e\s+hold\s+(?:that\s+)?(.{40,400}?)(?:\.|$)', 'holding'),
    (r'[Ww]e\s+hold\s+(?:only\s+)?(?:that\s+)?(.{40,400}?)(?:\.|$)', 'holding'),
    # Standard holding patterns
    (r'(?:we |the court )?hold(?:s|ing)? that (.+?)\.', 'holding'),
    (r'(?:we |the court )?conclude(?:s|d)? that (.+?)\.', 'holding'),
    (r'(?:we |the court )?find(?:s|ing)? that (.+?)\.', 'holding'),
    (r'(?:we |the court )?determine(?:s|d)? that (.+?)\.', 'holding'),
    (r'this (?:court|circuit) has (?:held|established|recognized) that (.+?)\.', 'holding'),
    (r'(?:the court|we) explained that (.+?)\.', 'holding'),

    # ===== ENUMERATED HOLDINGS (First, Second, Third) =====
    (r'[Ff]irst,?\s+(.{30,300}?)(?:\.|[Ss]econd)', 'holding'),
    (r'[Ss]econd,?\s+(.{30,300}?)(?:\.|[Tt]hird)', 'holding'),
    (r'[Tt]hird,?\s+(.{30,300}?)(?:\.|[Ff]ourth|\n\n)', 'holding'),
    (r'[Ff]ourth,?\s+(.{30,300}?)(?:\.|\n\n)', 'holding'),
    # "establishes the following" pattern
    (r'(?:establishes?|sets?\s+forth|provides?)\s+the\s+following\s+(?:prescriptions?|rules?|standards?|requirements?)[:\s]+(.{50,600}?)(?:\n\n|$)', 'holding'),

    # ===== STANDARDS AND TESTS =====
    (r'(?:the )?(?:applicable |governing )?(?:standard|test) (?:is|requires|provides) (.+?)\.', 'standard'),
    (r'(?:under|applying) (?:the )?(.+?) (?:test|standard|analysis|framework)', 'standard'),
    (r'(?:the )?(?:general |well[- ]established )?rule is that (.+?)\.', 'holding'),
    (r'(?:it is )?(?:well[- ])?(?:established|settled) (?:law )?that (.+?)\.', 'holding'),
    # Survival patterns (common in motions to dismiss)
    (r'(?:complaint|claim)\s+will\s+survive\s+(?:only\s+)?if\s+(.{40,400}?)(?:\.|$)', 'standard'),
    (r'(?:to\s+)?survive\s+(?:a\s+)?(?:motion\s+to\s+dismiss|12\(b\)\(6\)),?\s+(.{40,400}?)(?:\.|$)', 'standard'),

    # ===== ELEMENTS =====
    (r'(?:the )?elements? (?:are|include|of .+? are|consist of):?\s*(.+?)\.', 'elements'),
    (r'(?:to (?:establish|prove|show|prevail on) .+?,? )?(?:a |the )?(?:plaintiff|party|claimant) must (?:show|prove|establish|demonstrate) (.+?)\.', 'elements'),
    (r'(?:a )?(?:claim|cause of action) (?:for|of) .+? requires (?:proof of |that )?(.+?)\.', 'elements'),
    (r'requires? (?:a )?(?:showing|proof|demonstration) (?:of |that )(.+?)\.', 'elements'),
    (r'liability (?:attaches|arises|exists) (?:only )?(?:when|where|if) (.+?)\.', 'elements'),
    (r'(?:the )?(?:plaintiff|party|claimant) bears the burden (?:of|to) (.+?)\.', 'elements'),

    # ===== FACTORS =====
    (r'(?:the )?(?:relevant )?factors? (?:include|are|to be considered):?\s*(.+?)\.', 'factors'),
    (r'(?:courts? |we )?(?:consider|look to|examine|weigh) (?:the following |several )?(.+?)\.', 'factors'),
    (r'(?:in determining|when assessing) .+?,? courts? (?:consider|examine) (.+?)\.', 'factors'),

    # ===== INFERENCE PATTERNS (PSLRA specific) =====
    (r'(?:a |an )?(?:strong |cogent |compelling )?inference (?:of .+? )?must be (.+?)\.', 'standard'),
    (r'(?:a |an )?(?:strong )?inference (?:of .+? )?(?:arises|exists|may be drawn) (?:when|where|if) (.+?)\.', 'standard'),
    (r'(?:to |in order to )(?:support|establish) (?:a |the )?inference (?:of .+?)?,? (.+?)\.', 'elements'),
    (r'inference\s+of\s+scienter\s+must\s+be\s+(.{30,300}?)(?:\.|$)', 'standard'),
]

# Compile patterns
COMPILED_RULE_PATTERNS = [(re.compile(p, re.IGNORECASE | re.MULTILINE), t) for p, t in RULE_PATTERNS]


# ============================================================================
# ARGUMENT EXTRACTION PATTERNS
# ============================================================================

ARGUMENT_PATTERNS = [
    # Explicit argument markers
    r'(?:plaintiff|defendant|appellant|appellee|petitioner|respondent) (?:argues?|contends?|asserts?|maintains?) (?:that )?(.+?)(?:\.|$)',
    r'(?:the )?argument (?:is|that) (.+?)(?:\.|$)',
    r'(?:we|the court) (?:agree|disagree|reject|accept) (?:that )?(.+?)(?:\.|$)',

    # Because/therefore reasoning
    r'because (.+?), (?:we|the court) (.+?)(?:\.|$)',
    r'therefore,? (.+?)(?:\.|$)',
    r'accordingly,? (.+?)(?:\.|$)',
    r'thus,? (.+?)(?:\.|$)',

    # Conclusion markers
    r'(?:we )?conclude (?:that )?(.+?)(?:\.|$)',
    r'(?:for (?:these|the foregoing) reasons?,? )(.+?)(?:\.|$)',

    # Brief heading patterns (Roman numerals, letters, numbers)
    r'^(?:I+|IV|V|VI+|IX|X+)\.?\s+(.+?)(?:\n|$)',
    r'^[A-Z]\.?\s+(.+?)(?:\n|$)',
    r'^\d+\.?\s+(.+?)(?:\n|$)',

    # Common claim structures
    r'(?:is|was|were|are) (?:entitled to|liable for|responsible for) (.+?)(?:\.|$)',
    r'(?:clearly|obviously|plainly|unambiguously) (.+?)(?:\.|$)',
    r'(?:must|should|shall) (?:be )?(.+?)(?:\.|$)',
    r'(?:failed|neglected|refused) to (.+?)(?:\.|$)',
    r'(?:violated|breached|broke) (?:the )?(.+?)(?:\.|$)',
]

COMPILED_ARG_PATTERNS = [re.compile(p, re.IGNORECASE | re.MULTILINE) for p in ARGUMENT_PATTERNS]


# ============================================================================
# FACT EXTRACTION PATTERNS
# ============================================================================

PROCEDURAL_PATTERNS = [
    r'(?:this )?(?:case|matter|action) (?:comes? before|is before|arises from) (.+?)(?:\.|$)',
    r'(?:plaintiff|defendant) (?:appeals?|moved?|filed) (.+?)(?:\.|$)',
    r'(?:the )?(?:district|trial|lower) court (?:held|ruled|found|granted|denied) (.+?)(?:\.|$)',
    r'(?:we )?(?:review|consider|address) (.+?)(?:\.|$)',
    r'(?:on|following) (?:appeal|review),? (.+?)(?:\.|$)',
]

FACT_INDICATORS = [
    r'(?:the )?facts (?:are|show|indicate|establish) (?:that )?(.+?)(?:\.|$)',
    r'(?:in|on|during) (?:january|february|march|april|may|june|july|august|september|october|november|december) \d{1,2},? \d{4},? (.+?)(?:\.|$)',
    r'(?:plaintiff|defendant|appellant|appellee) (.+?)(?:\.|$)',
]

COMPILED_PROC_PATTERNS = [re.compile(p, re.IGNORECASE) for p in PROCEDURAL_PATTERNS]
COMPILED_FACT_PATTERNS = [re.compile(p, re.IGNORECASE) for p in FACT_INDICATORS]


# ============================================================================
# RULE SYNTHESIS FUNCTIONS
# ============================================================================

def extract_rules_from_text(text: str, source_case: str = "Unknown") -> List[LegalRule]:
    """
    Extract black letter rules from case text.

    Args:
        text: Case opinion or brief text
        source_case: Name of the source case

    Returns:
        List of LegalRule objects
    """
    rules = []

    for pattern, rule_type in COMPILED_RULE_PATTERNS:
        matches = pattern.finditer(text)
        for match in matches:
            rule_text = match.group(1).strip()

            # Skip very short matches (likely false positives)
            if len(rule_text) < 20:
                continue

            # Extract surrounding context
            start = max(0, match.start() - 100)
            end = min(len(text), match.end() + 100)
            context = text[start:end].strip()

            # Extract elements if it's an elements-type rule
            elements = []
            if rule_type == 'elements':
                # Look for numbered or lettered lists
                element_matches = re.findall(r'(?:\d+\)|[a-z]\)|•|-)\s*([^;,\n]+)', rule_text)
                if element_matches:
                    elements = [e.strip() for e in element_matches]

            rule = LegalRule(
                text=rule_text,
                source_case=source_case,
                rule_type=rule_type,
                elements=elements,
                context=context
            )
            rules.append(rule)

    # FALLBACK: If no rules found, extract sentences with legal keywords
    if not rules:
        rules = _extract_rules_fallback(text, source_case)

    return rules


def _extract_rules_fallback(text: str, source_case: str) -> List[LegalRule]:
    """
    Fallback extraction: find sentences containing legal rule indicators.
    """
    rules = []
    seen = set()

    # Keywords indicating a rule
    indicators = re.compile(
        r'\b(?:must|shall|requires?|burden|elements?|factors?|standard|test|'
        r'prove|establish|demonstrate|show|liability|inference|plaintiff|defendant)\b',
        re.IGNORECASE
    )

    # Split into sentences
    sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', text)

    for sentence in sentences:
        sentence = sentence.strip()

        # Must be substantial but not too long
        if len(sentence) < 50 or len(sentence) > 400:
            continue

        # Must contain a rule indicator
        if not indicators.search(sentence):
            continue

        # Avoid duplicates
        normalized = ' '.join(sentence.lower().split())
        if normalized in seen:
            continue
        seen.add(normalized)

        # Skip citations and headings
        if re.match(r'^[\d\s]+[A-Z]\.', sentence) or sentence.isupper():
            continue

        # Determine type
        if re.search(r'\belements?\b', sentence, re.IGNORECASE):
            rule_type = 'elements'
        elif re.search(r'\bfactors?\b', sentence, re.IGNORECASE):
            rule_type = 'factors'
        else:
            rule_type = 'holding'

        rules.append(LegalRule(
            text=sentence,
            source_case=source_case,
            rule_type=rule_type,
            elements=[],
            context=""
        ))

        if len(rules) >= 8:
            break

    return rules


def synthesize_rules(case_texts: Dict[str, str]) -> Dict[str, List[LegalRule]]:
    """
    Synthesize rules from multiple cases.

    Args:
        case_texts: Dictionary mapping case name -> case text

    Returns:
        Dictionary mapping rule type -> list of rules
    """
    all_rules = defaultdict(list)

    for case_name, text in case_texts.items():
        rules = extract_rules_from_text(text, case_name)
        for rule in rules:
            all_rules[rule.rule_type].append(rule)

    return dict(all_rules)


def format_synthesized_rules(rules_by_type: Dict[str, List[LegalRule]]) -> str:
    """Format synthesized rules for output."""
    output = []
    output.append("SYNTHESIZED LEGAL RULES")
    output.append("=" * 70)

    for rule_type, rules in rules_by_type.items():
        # Pluralize correctly - don't add S if already ends in S
        type_name = rule_type.upper()
        if not type_name.endswith('S'):
            type_name += 'S'
        output.append(f"\n## {type_name}")
        output.append("-" * 70)

        for i, rule in enumerate(rules, 1):
            output.append(f"\n{i}. From {rule.source_case}:")
            output.append(f"   \"{rule.text}\"")

            if rule.elements:
                output.append("   Elements:")
                for j, element in enumerate(rule.elements, 1):
                    output.append(f"     {j}. {element}")

    return '\n'.join(output)


# ============================================================================
# ARGUMENT EXTRACTION FUNCTIONS
# ============================================================================

def extract_arguments_from_text(text: str, citations: Dict[str, Dict] = None) -> List[Argument]:
    """
    Extract discrete legal arguments from text.

    Args:
        text: Brief or opinion text
        citations: Optional dictionary of citations found in text

    Returns:
        List of Argument objects
    """
    arguments = []
    citations = citations or {}

    for pattern in COMPILED_ARG_PATTERNS:
        matches = pattern.finditer(text)
        for match in matches:
            claim = match.group(1).strip() if match.lastindex else match.group(0).strip()

            # Skip very short matches
            if len(claim) < 30:
                continue

            # Find supporting citations near this argument
            start = max(0, match.start() - 200)
            end = min(len(text), match.end() + 200)
            context = text[start:end]

            supporting = []
            for citation in citations.keys():
                if citation in context:
                    supporting.append(citation)

            # Determine strength based on number of citations
            strength = min(5, len(supporting) + 1)

            # Determine argument type
            arg_type = "main"
            if re.search(r'alternatively|in the alternative', claim, re.IGNORECASE):
                arg_type = "alternative"
            elif re.search(r'however|but|although|counter', claim, re.IGNORECASE):
                arg_type = "counter"

            arg = Argument(
                claim=claim,
                supporting_authority=supporting,
                reasoning=context,
                strength=strength,
                argument_type=arg_type
            )
            arguments.append(arg)

    return arguments


def format_arguments(arguments: List[Argument]) -> str:
    """Format extracted arguments for output."""
    output = []
    output.append("EXTRACTED ARGUMENTS")
    output.append("=" * 70)
    output.append(f"\nTotal arguments found: {len(arguments)}\n")

    # Group by type
    by_type = defaultdict(list)
    for arg in arguments:
        by_type[arg.argument_type].append(arg)

    for arg_type in ["main", "alternative", "counter"]:
        args = by_type.get(arg_type, [])
        if not args:
            continue

        output.append(f"\n## {arg_type.upper()} ARGUMENTS")
        output.append("-" * 70)

        for i, arg in enumerate(args, 1):
            output.append(f"\n{i}. {arg.claim}")
            output.append(f"   Strength: {'*' * arg.strength} ({arg.strength}/5)")
            if arg.supporting_authority:
                output.append(f"   Supporting Authority:")
                for cite in arg.supporting_authority:
                    output.append(f"     - {cite}")

    return '\n'.join(output)


# ============================================================================
# FACT PATTERN EXTRACTION FUNCTIONS
# ============================================================================

def extract_fact_pattern(text: str) -> FactPattern:
    """
    Extract fact pattern from case or brief.

    Args:
        text: Case opinion or brief text

    Returns:
        FactPattern object
    """
    # Extract procedural posture
    procedural = ""
    for pattern in COMPILED_PROC_PATTERNS:
        match = pattern.search(text)
        if match:
            procedural = match.group(0).strip()
            break

    # Extract facts
    key_facts = []
    background_facts = []

    # Look for facts section
    facts_section = re.search(
        r'(?:FACTS?|BACKGROUND|STATEMENT OF FACTS?)\s*\n(.+?)(?=\n(?:ISSUE|DISCUSSION|ANALYSIS|ARGUMENT|HOLDING)|\Z)',
        text,
        re.IGNORECASE | re.DOTALL
    )

    if facts_section:
        facts_text = facts_section.group(1)
        sentences = re.split(r'(?<=[.!?])\s+', facts_text)

        for sentence in sentences:
            sentence = sentence.strip()
            if len(sentence) < 20:
                continue

            # Legally significant facts often contain legal terms
            if re.search(r'contract|agreement|breach|damage|injur|negligen|intent|knew|should have known',
                        sentence, re.IGNORECASE):
                key_facts.append(sentence)
            else:
                background_facts.append(sentence)

    # Extract party names
    parties = {}
    party_patterns = [
        (r'(?:plaintiff|petitioner|appellant)[,\s]+([A-Z][a-zA-Z\s,\.]+?)(?:,|\.|v\.)', 'plaintiff'),
        (r'(?:defendant|respondent|appellee)[,\s]+([A-Z][a-zA-Z\s,\.]+?)(?:,|\.|v\.)', 'defendant'),
    ]

    for pattern, role in party_patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            parties[role] = match.group(1).strip()

    return FactPattern(
        procedural_posture=procedural,
        key_facts=key_facts[:10],  # Limit to top 10
        background_facts=background_facts[:5],  # Limit to 5
        legally_significant_facts=key_facts[:5],  # Top 5 are most significant
        parties=parties
    )


def format_fact_pattern(facts: FactPattern) -> str:
    """Format fact pattern for output."""
    output = []
    output.append("FACT PATTERN ANALYSIS")
    output.append("=" * 70)

    if facts.parties:
        output.append("\n## PARTIES")
        for role, name in facts.parties.items():
            output.append(f"  {role.capitalize()}: {name}")

    if facts.procedural_posture:
        output.append("\n## PROCEDURAL POSTURE")
        output.append(f"  {facts.procedural_posture}")

    if facts.legally_significant_facts:
        output.append("\n## LEGALLY SIGNIFICANT FACTS")
        for i, fact in enumerate(facts.legally_significant_facts, 1):
            output.append(f"  {i}. {fact}")

    if facts.key_facts:
        output.append("\n## KEY FACTS")
        for i, fact in enumerate(facts.key_facts, 1):
            output.append(f"  {i}. {fact}")

    if facts.background_facts:
        output.append("\n## BACKGROUND FACTS")
        for i, fact in enumerate(facts.background_facts, 1):
            output.append(f"  {i}. {fact}")

    return '\n'.join(output)


# ============================================================================
# CASE COMPARISON FUNCTIONS
# ============================================================================

def compare_cases(case_a_text: str, case_b_text: str,
                  case_a_name: str = "Case A", case_b_name: str = "Case B") -> CaseComparison:
    """
    Compare two cases side-by-side for analogies and distinctions.

    Args:
        case_a_text: Text of first case
        case_b_text: Text of second case
        case_a_name: Name of first case
        case_b_name: Name of second case

    Returns:
        CaseComparison object with similarities, differences, and draft language
    """
    # Use improved extraction methods
    from enhanced_rules import extract_enhanced_rules
    from brief_synthesis import extract_structured_case_facts

    # Extract structured facts
    facts_a = extract_structured_case_facts(case_a_text)
    facts_b = extract_structured_case_facts(case_b_text)

    # Extract rules using improved extraction
    rules_a = extract_enhanced_rules(case_a_text, case_a_name)
    rules_b = extract_enhanced_rules(case_b_text, case_b_name)

    similarities = []
    differences = []
    distinguishing_factors = []
    suggested_language = []

    # Compare outcomes
    if facts_a.outcome and facts_b.outcome:
        outcome_sim = _text_similarity(facts_a.outcome, facts_b.outcome)
        if outcome_sim > 0.4:
            similarities.append(f"Similar outcome: Both cases {facts_a.outcome[:60]}")
        else:
            differences.append(f"Different outcomes: {case_a_name} ({facts_a.outcome[:40]}...) vs {case_b_name} ({facts_b.outcome[:40]}...)")

    # Compare rules - find overlapping and unique rules
    rules_a_texts = [r.text for r in rules_a]
    rules_b_texts = [r.text for r in rules_b]

    # Track which rules match
    matched_a = set()
    matched_b = set()

    for i, rule_a in enumerate(rules_a):
        for j, rule_b in enumerate(rules_b):
            if j in matched_b:
                continue
            sim = _text_similarity(rule_a.text, rule_b.text)
            if sim > 0.35:  # Lower threshold for legal text
                similarities.append(f"Both apply: \"{rule_a.text[:100]}...\"" if len(rule_a.text) > 100 else f"Both apply: \"{rule_a.text}\"")
                matched_a.add(i)
                matched_b.add(j)
                break

    # Unique rules
    for i, rule in enumerate(rules_a):
        if i not in matched_a:
            differences.append(f"{case_a_name} rule: \"{rule.text[:80]}...\"" if len(rule.text) > 80 else f"{case_a_name} rule: \"{rule.text}\"")

    for j, rule in enumerate(rules_b):
        if j not in matched_b:
            differences.append(f"{case_b_name} rule: \"{rule.text[:80]}...\"" if len(rule.text) > 80 else f"{case_b_name} rule: \"{rule.text}\"")

    # Compare determinative facts - these are key for analogies/distinctions
    for fact_a in facts_a.determinative_facts[:3]:
        found_similar = False
        for fact_b in facts_b.determinative_facts:
            if _text_similarity(fact_a, fact_b) > 0.3:
                similarities.append(f"Similar determinative fact: {fact_a[:80]}")
                found_similar = True
                break
        if not found_similar:
            distinguishing_factors.append(f"{case_a_name}: {fact_a[:100]}")

    for fact_b in facts_b.determinative_facts[:3]:
        if not any(_text_similarity(fact_b, fa) > 0.3 for fa in facts_a.determinative_facts):
            distinguishing_factors.append(f"{case_b_name}: {fact_b[:100]}")

    # Compare operative facts for additional context
    all_facts_a = facts_a.operative_facts + facts_a.undisputed_facts
    all_facts_b = facts_b.operative_facts + facts_b.undisputed_facts

    for fact_a in all_facts_a[:5]:
        for fact_b in all_facts_b[:5]:
            if _text_similarity(fact_a, fact_b) > 0.4:
                if f"Similar fact" not in str(similarities):  # Avoid too many
                    similarities.append(f"Similar fact pattern: {fact_a[:60]}...")
                break

    # Generate suggested distinction/analogy language
    if similarities:
        suggested_language.append(f"\n## ANALOGY LANGUAGE (if {case_a_name} supports your position):")
        suggested_language.append(f"Like {case_a_name}, where {facts_a.outcome[:50] if facts_a.outcome else '[outcome]'},")
        suggested_language.append(f"here [describe your similar facts].")

    if distinguishing_factors:
        suggested_language.append(f"\n## DISTINCTION LANGUAGE (if {case_b_name} is adverse):")
        suggested_language.append(f"Unlike {case_b_name}, where {distinguishing_factors[0][:60]}...,")
        suggested_language.append(f"here [describe your distinguishing facts].")

    return CaseComparison(
        similarities=similarities if similarities else ["No clear similarities found - cases may address different issues"],
        differences=differences if differences else ["No clear differences found"],
        distinguishing_factors=distinguishing_factors if distinguishing_factors else ["Insufficient facts extracted for comparison"],
        suggested_language=suggested_language if suggested_language else ["Provide more detailed case text for draft language"]
    )


def _text_similarity(text_a: str, text_b: str) -> float:
    """Simple word overlap similarity."""
    words_a = set(text_a.lower().split())
    words_b = set(text_b.lower().split())

    if not words_a or not words_b:
        return 0.0

    intersection = words_a & words_b
    union = words_a | words_b

    return len(intersection) / len(union)


def format_case_comparison(comparison: CaseComparison) -> str:
    """Format case comparison for output."""
    output = []
    output.append("CASE COMPARISON ANALYSIS")
    output.append("=" * 70)

    if comparison.similarities:
        output.append("\n## SIMILARITIES")
        for sim in comparison.similarities:
            output.append(f"  - {sim}")

    if comparison.differences:
        output.append("\n## DIFFERENCES")
        for diff in comparison.differences:
            output.append(f"  - {diff}")

    if comparison.distinguishing_factors:
        output.append("\n## KEY DISTINGUISHING FACTORS")
        for factor in comparison.distinguishing_factors:
            output.append(f"  - {factor}")

    if comparison.suggested_language:
        output.append("\n## SUGGESTED DISTINCTION LANGUAGE")
        output.append("-" * 70)
        for line in comparison.suggested_language:
            output.append(line)

    return '\n'.join(output)


# ============================================================================
# MEMO DRAFT GENERATOR
# ============================================================================

def generate_memo_draft(issue: str, facts: str, case_texts: Dict[str, str]) -> MemoDraft:
    """
    Generate a draft legal memo structure.

    Args:
        issue: The legal issue/question
        facts: Your client's facts
        case_texts: Dictionary mapping case name -> case text

    Returns:
        MemoDraft object
    """
    # Synthesize rules from cases
    synthesized = synthesize_rules(case_texts)

    # Extract main holdings
    holdings = synthesized.get('holding', [])

    # Generate Question Presented
    question_presented = f"Whether {issue}"
    if not question_presented.endswith('?'):
        question_presented += '?'

    # Generate Brief Answer
    brief_answer = "Probably [yes/no]. "
    if holdings:
        brief_answer += f"Under {holdings[0].source_case}, {holdings[0].text[:100]}..."

    # Generate Facts Template
    facts_template = f"""STATEMENT OF FACTS

[Client name] seeks advice regarding {issue}.

{facts}

[Additional relevant facts from client interview/documents]
"""

    # Generate IRAC Discussion
    discussion = {
        'issue': f"The issue is {issue}.",
        'rule': "",
        'application': "",
        'conclusion': ""
    }

    # Build rule section from synthesized rules
    if holdings:
        rule_text = "The applicable legal standard is established by the following authorities:\n\n"
        for rule in holdings[:3]:
            rule_text += f"In {rule.source_case}, the court held that {rule.text}\n\n"
        discussion['rule'] = rule_text

    # Application template
    discussion['application'] = f"""Applying these principles to the present facts:

[For each element/factor of the rule, apply to client's facts]

1. [First element]: Here, [apply facts to element]...

2. [Second element]: Similarly, [apply facts to element]...

[Continue for each element]
"""

    # Conclusion template
    discussion['conclusion'] = f"""Based on the foregoing analysis, [client] [likely will/will not] prevail on the {issue} claim because [summarize key reasons]."""

    # Final conclusion
    conclusion = f"For the reasons stated above, [summary recommendation]."

    return MemoDraft(
        question_presented=question_presented,
        brief_answer=brief_answer,
        facts_template=facts_template,
        discussion_irac=discussion,
        conclusion=conclusion
    )


def format_memo_draft(memo: MemoDraft) -> str:
    """Format memo draft for output."""
    output = []
    output.append("LEGAL MEMORANDUM DRAFT")
    output.append("=" * 70)

    output.append("\n## QUESTION PRESENTED")
    output.append("-" * 70)
    output.append(memo.question_presented)

    output.append("\n\n## BRIEF ANSWER")
    output.append("-" * 70)
    output.append(memo.brief_answer)

    output.append("\n\n## STATEMENT OF FACTS")
    output.append("-" * 70)
    output.append(memo.facts_template)

    output.append("\n\n## DISCUSSION")
    output.append("-" * 70)

    output.append("\n### Issue")
    output.append(memo.discussion_irac['issue'])

    output.append("\n### Rule")
    output.append(memo.discussion_irac['rule'])

    output.append("\n### Application")
    output.append(memo.discussion_irac['application'])

    output.append("\n### Conclusion")
    output.append(memo.discussion_irac['conclusion'])

    output.append("\n\n## CONCLUSION")
    output.append("-" * 70)
    output.append(memo.conclusion)

    return '\n'.join(output)


# ============================================================================
# CONVENIENCE FUNCTIONS FOR MCP
# ============================================================================

def analyze_for_rules(text: str, case_name: str = "Unknown") -> str:
    """Extract and format rules from a single case."""
    rules = extract_rules_from_text(text, case_name)
    return format_synthesized_rules({r.rule_type: [r] for r in rules})


def analyze_for_arguments(text: str, citations: Dict = None) -> str:
    """Extract and format arguments from text."""
    arguments = extract_arguments_from_text(text, citations)
    return format_arguments(arguments)


def analyze_for_facts(text: str) -> str:
    """Extract and format fact pattern from text."""
    facts = extract_fact_pattern(text)
    return format_fact_pattern(facts)


def compare_two_cases(case_a_text: str, case_b_text: str,
                      case_a_name: str = "Case A", case_b_name: str = "Case B") -> str:
    """Compare two cases and return formatted output."""
    comparison = compare_cases(case_a_text, case_b_text, case_a_name, case_b_name)
    return format_case_comparison(comparison)


def generate_memo(issue: str, facts: str, case_texts: Dict[str, str]) -> str:
    """Generate and format a memo draft."""
    memo = generate_memo_draft(issue, facts, case_texts)
    return format_memo_draft(memo)


# ============================================================================
# OPPOSING BRIEF ANALYZER
# ============================================================================

@dataclass
class CounterArgument:
    """Specific counter-argument with draft language."""
    opponent_argument: str
    counter_type: str  # "factual", "legal", "procedural", "authority"
    counter_strategy: str
    draft_language: str
    supporting_authority: Optional[str] = None


@dataclass
class OpposingBriefAnalysis:
    """Complete analysis of opposing counsel's brief"""
    arguments: List[Argument]
    citations_used: Dict[str, Dict]
    weak_citations: List[str]  # Citations that might be distinguishable
    counter_argument_suggestions: List[str]
    fact_disputes: List[str]
    authority_gaps: List[str]  # Topics where opponent has weak authority
    theory_of_case: str = ""  # Their narrative
    disputed_facts: List[str] = field(default_factory=list)
    undisputed_facts: List[str] = field(default_factory=list)
    procedural_arguments: List[str] = field(default_factory=list)
    specific_counters: List[CounterArgument] = field(default_factory=list)


def analyze_opposing_brief(brief_text: str, your_facts: str = "") -> OpposingBriefAnalysis:
    """
    Analyze opposing counsel's brief and generate counter-argument strategy.

    Enhanced features:
    - Extracts their theory of the case (narrative)
    - Identifies disputed vs. undisputed facts
    - Spots procedural arguments
    - Generates specific counter-argument language

    Args:
        brief_text: The opposing brief text
        your_facts: Your client's facts (optional, for comparison)

    Returns:
        OpposingBriefAnalysis object
    """
    from text_extraction_core import extract_citations_from_text

    # Extract their arguments
    arguments = extract_arguments_from_text(brief_text)

    # Extract their citations
    citations = extract_citations_from_text(brief_text)

    # NEW: Extract their theory of the case
    theory_of_case = _extract_theory_of_case(brief_text)

    # Identify potentially weak citations
    weak_citations = []
    for citation in citations.keys():
        # Check for old cases (pre-2000 might be distinguishable)
        year_match = re.search(r'\((\d{4})\)', citation)
        if year_match and int(year_match.group(1)) < 2000:
            weak_citations.append(f"{citation} - OLD CASE (may be distinguishable or superseded)")

        # Check for non-binding authority
        if re.search(r'F\.\s*Supp|Bankr\.|App\'x', citation):
            weak_citations.append(f"{citation} - LOWER COURT/UNPUBLISHED (persuasive only)")

    # Generate counter-argument suggestions (legacy)
    counter_suggestions = []
    for arg in arguments:
        if arg.argument_type == "main":
            counter_suggestions.append(
                f"Counter to: \"{arg.claim[:80]}...\"\n"
                f"  Suggested approach: Challenge the [legal standard/factual basis/authority]"
            )
            if arg.supporting_authority:
                counter_suggestions.append(
                    f"  Distinguish: {', '.join(arg.supporting_authority[:2])}"
                )

    # NEW: Generate SPECIFIC counter-arguments with draft language
    specific_counters = _generate_specific_counters(arguments, brief_text, your_facts)

    # Identify fact disputes if your facts provided
    fact_disputes = []
    disputed_facts = []
    undisputed_facts = []
    if your_facts:
        opp_facts = extract_fact_pattern(brief_text)
        your_facts_lower = your_facts.lower()

        for opp_fact in opp_facts.key_facts[:10]:
            opp_fact_lower = opp_fact.lower()
            # Check if fact matches your version
            matching_words = sum(1 for word in opp_fact_lower.split() if word in your_facts_lower)
            total_words = len(opp_fact_lower.split())

            if matching_words / max(total_words, 1) > 0.6:
                undisputed_facts.append(opp_fact)
            else:
                disputed_facts.append(opp_fact)
                fact_disputes.append(f"Opponent claims: \"{opp_fact[:100]}...\" - VERIFY/DISPUTE")

    # Identify authority gaps
    authority_gaps = []
    for arg in arguments:
        if len(arg.supporting_authority) < 2:
            authority_gaps.append(
                f"WEAK SUPPORT: \"{arg.claim[:60]}...\" has only {len(arg.supporting_authority)} citation(s)"
            )

    # NEW: Identify procedural arguments
    procedural_arguments = _extract_procedural_arguments(brief_text)

    return OpposingBriefAnalysis(
        arguments=arguments,
        citations_used=citations,
        weak_citations=weak_citations,
        counter_argument_suggestions=counter_suggestions,
        fact_disputes=fact_disputes,
        authority_gaps=authority_gaps,
        theory_of_case=theory_of_case,
        disputed_facts=disputed_facts,
        undisputed_facts=undisputed_facts,
        procedural_arguments=procedural_arguments,
        specific_counters=specific_counters
    )


def _extract_theory_of_case(brief_text: str) -> str:
    """Extract opponent's theory/narrative of the case."""
    text_lower = brief_text.lower()

    # Look for theory statements
    theory_patterns = [
        r'(?:this\s+(?:case|matter)\s+(?:is\s+about|concerns|involves)\s+)([^.]+\.)',
        r'(?:the\s+(?:evidence|facts|record)\s+(?:shows?|demonstrates?|establishes?)\s+that\s+)([^.]+\.)',
        r'(?:(?:plaintiff|defendant)\s+(?:seeks?|claims?|alleges?)\s+)([^.]+\.)',
        r'(?:at\s+its\s+core,?\s+)([^.]+\.)',
        r'(?:the\s+central\s+(?:issue|question)\s+is\s+)([^.]+\.)',
    ]

    theories = []
    for pattern in theory_patterns:
        matches = re.findall(pattern, brief_text, re.IGNORECASE)
        theories.extend(matches[:2])

    # Also check introduction/summary sections
    intro_section = re.search(
        r'(?:INTRODUCTION|SUMMARY|PRELIMINARY\s+STATEMENT)[:\s]+(.+?)(?:FACTS|ARGUMENT|ISSUE)',
        brief_text, re.IGNORECASE | re.DOTALL
    )
    if intro_section:
        # Get first few sentences
        intro_sentences = re.split(r'(?<=[.!?])\s+', intro_section.group(1))
        theories.extend(s.strip() for s in intro_sentences[:3] if len(s.strip()) > 30)

    if theories:
        return "OPPONENT'S NARRATIVE: " + " ".join(theories[:2])[:300]
    return "Could not identify clear theory of case"


def _extract_procedural_arguments(brief_text: str) -> List[str]:
    """Extract procedural arguments from brief."""
    procedural_args = []

    procedural_patterns = [
        (r'(?:motion\s+to\s+dismiss|12\(b\)\(6\)|failure\s+to\s+state\s+a\s+claim)[^.]+\.', 'MTD'),
        (r'(?:summary\s+judgment|rule\s+56|no\s+genuine\s+(?:dispute|issue))[^.]+\.', 'MSJ'),
        (r'(?:standing|lack\s+of\s+standing|article\s+iii)[^.]+\.', 'Standing'),
        (r'(?:statute\s+of\s+limitations?|time.?barred)[^.]+\.', 'SOL'),
        (r'(?:jurisdiction|lack\s+of\s+jurisdiction|subject.?matter)[^.]+\.', 'Jurisdiction'),
        (r'(?:exhaustion|failed\s+to\s+exhaust|administrative\s+remedies)[^.]+\.', 'Exhaustion'),
        (r'(?:ripeness|not\s+(?:yet\s+)?ripe|premature)[^.]+\.', 'Ripeness'),
        (r'(?:mootness|moot|no\s+longer\s+a\s+(?:live\s+)?controversy)[^.]+\.', 'Mootness'),
    ]

    for pattern, arg_type in procedural_patterns:
        matches = re.findall(pattern, brief_text, re.IGNORECASE)
        for match in matches[:2]:
            procedural_args.append(f"[{arg_type}] {match[:150]}")

    return procedural_args[:5]


def _generate_specific_counters(
    arguments: List[Argument],
    brief_text: str,
    your_facts: str
) -> List[CounterArgument]:
    """Generate specific counter-arguments with draft language."""
    counters = []

    for arg in arguments[:5]:  # Limit to main arguments
        if arg.argument_type != "main":
            continue

        claim_lower = arg.claim.lower()

        # Determine counter type and generate draft language
        if any(word in claim_lower for word in ['undisputed', 'no genuine', 'clear', 'obvious']):
            # Factual counter
            counters.append(CounterArgument(
                opponent_argument=arg.claim[:100],
                counter_type="factual",
                counter_strategy="Challenge factual premise",
                draft_language=(
                    f"Contrary to [Defendant/Plaintiff]'s assertion, the record demonstrates genuine disputes "
                    f"of material fact. Specifically, [cite specific contradicting evidence]. See [Record cite]. "
                    f"This disputed evidence precludes the relief [Defendant/Plaintiff] seeks."
                ),
                supporting_authority="[Insert case re: genuine disputes preclude summary judgment]"
            ))

        elif any(word in claim_lower for word in ['established', 'settled', 'binding', 'controlling']):
            # Legal authority counter
            auth_list = ', '.join(arg.supporting_authority[:2]) if arg.supporting_authority else '[their case]'
            counters.append(CounterArgument(
                opponent_argument=arg.claim[:100],
                counter_type="legal",
                counter_strategy="Distinguish cited authority",
                draft_language=(
                    f"[Defendant/Plaintiff] relies on {auth_list}, but that case is distinguishable. "
                    f"There, the court addressed [precedent's factual situation]. Here, by contrast, "
                    f"[your distinguishing facts]. This critical distinction renders [their case] inapplicable."
                ),
                supporting_authority="[Insert your supporting case]"
            ))

        elif any(word in claim_lower for word in ['failed', 'cannot', 'did not', 'lacks']):
            # Affirmative counter
            counters.append(CounterArgument(
                opponent_argument=arg.claim[:100],
                counter_type="affirmative",
                counter_strategy="Affirmatively prove the element",
                draft_language=(
                    f"[Plaintiff/Defendant] has, in fact, established [the disputed element]. "
                    f"The evidence shows [specific facts]. See [Record cite]. "
                    f"This evidence satisfies the [legal standard] standard under [case law]."
                ),
                supporting_authority="[Insert case establishing standard]"
            ))

        else:
            # General counter
            counters.append(CounterArgument(
                opponent_argument=arg.claim[:100],
                counter_type="general",
                counter_strategy="Address directly",
                draft_language=(
                    f"[Defendant/Plaintiff]'s argument fails for [X] reasons. First, [reason 1]. "
                    f"Second, [reason 2]. The authorities [Defendant/Plaintiff] cites do not support "
                    f"the proposition advanced because [distinguish/limit]."
                ),
                supporting_authority="[Insert your authority]"
            ))

    return counters


def format_opposing_brief_analysis(analysis: OpposingBriefAnalysis) -> str:
    """Format opposing brief analysis for output."""
    output = []
    output.append("OPPOSING BRIEF ANALYSIS")
    output.append("=" * 70)
    output.append("Strategic Analysis for Counter-Arguments")
    output.append("=" * 70)

    # Summary
    output.append(f"\n## SUMMARY")
    output.append(f"  Total arguments identified: {len(analysis.arguments)}")
    output.append(f"  Total citations used: {len(analysis.citations_used)}")
    output.append(f"  Potentially weak citations: {len(analysis.weak_citations)}")
    output.append(f"  Authority gaps identified: {len(analysis.authority_gaps)}")

    # NEW: Theory of the case
    if analysis.theory_of_case:
        output.append(f"\n## OPPONENT'S THEORY OF THE CASE")
        output.append("-" * 70)
        output.append(f"  {analysis.theory_of_case}")

    # Their arguments
    output.append(f"\n## OPPONENT'S ARGUMENTS")
    output.append("-" * 70)
    for i, arg in enumerate(analysis.arguments[:10], 1):
        output.append(f"\n{i}. [{arg.argument_type.upper()}] {arg.claim[:150]}")
        if arg.supporting_authority:
            output.append(f"   Supporting: {', '.join(arg.supporting_authority[:3])}")
        output.append(f"   Strength: {'*' * arg.strength}/5")

    # NEW: Procedural arguments
    if analysis.procedural_arguments:
        output.append(f"\n## PROCEDURAL ARGUMENTS TO ADDRESS")
        output.append("-" * 70)
        for proc_arg in analysis.procedural_arguments:
            output.append(f"  {proc_arg}")

    # Weak citations to attack
    if analysis.weak_citations:
        output.append(f"\n## WEAK CITATIONS TO ATTACK")
        output.append("-" * 70)
        for cite in analysis.weak_citations[:10]:
            output.append(f"  - {cite}")

    # NEW: Specific counter-arguments with draft language
    if analysis.specific_counters:
        output.append(f"\n## SPECIFIC COUNTER-ARGUMENTS (DRAFT LANGUAGE)")
        output.append("-" * 70)
        for i, counter in enumerate(analysis.specific_counters, 1):
            output.append(f"\n[{i}] [{counter.counter_type.upper()}] Responding to: \"{counter.opponent_argument[:60]}...\"")
            output.append(f"    Strategy: {counter.counter_strategy}")
            output.append(f"    DRAFT:")
            output.append(f"      \"{counter.draft_language}\"")
            if counter.supporting_authority:
                output.append(f"    Cite: {counter.supporting_authority}")

    # Legacy counter-argument suggestions
    if analysis.counter_argument_suggestions and not analysis.specific_counters:
        output.append(f"\n## COUNTER-ARGUMENT STRATEGY")
        output.append("-" * 70)
        for suggestion in analysis.counter_argument_suggestions[:10]:
            output.append(f"\n{suggestion}")

    # Authority gaps
    if analysis.authority_gaps:
        output.append(f"\n## AUTHORITY GAPS (EXPLOIT THESE)")
        output.append("-" * 70)
        for gap in analysis.authority_gaps:
            output.append(f"  - {gap}")

    # NEW: Disputed vs. Undisputed facts
    if analysis.disputed_facts or analysis.undisputed_facts:
        output.append(f"\n## FACT CLASSIFICATION")
        output.append("-" * 70)
        if analysis.undisputed_facts:
            output.append(f"  UNDISPUTED ({len(analysis.undisputed_facts)}):")
            for fact in analysis.undisputed_facts[:3]:
                output.append(f"    [.] {fact[:80]}...")
        if analysis.disputed_facts:
            output.append(f"  DISPUTED ({len(analysis.disputed_facts)}):")
            for fact in analysis.disputed_facts[:3]:
                output.append(f"    [!] {fact[:80]}...")

    # Legacy fact disputes
    if analysis.fact_disputes and not analysis.disputed_facts:
        output.append(f"\n## POTENTIAL FACT DISPUTES")
        output.append("-" * 70)
        for dispute in analysis.fact_disputes:
            output.append(f"  - {dispute}")

    return '\n'.join(output)


def analyze_opposing(brief_text: str, your_facts: str = "") -> str:
    """Analyze opposing brief and return formatted output."""
    analysis = analyze_opposing_brief(brief_text, your_facts)
    return format_opposing_brief_analysis(analysis)
