"""
Case Brief Generator and Argument Analysis Tools
AI-powered extraction of legal reasoning and argument structure
"""

import re
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from enum import Enum


# ============================================================================
# CASE BRIEF COMPONENTS
# ============================================================================

@dataclass
class CaseBrief:
    """Structured case brief"""
    case_name: str
    citation: str
    court: str
    year: str
    facts: List[str]
    procedural_history: List[str]
    issues: List[str]
    holdings: List[str]
    reasoning: List[str]
    rule: List[str]
    dissents: List[str]
    concurrences: List[str]


class SectionType(Enum):
    """Types of sections in a legal opinion"""
    FACTS = "facts"
    PROCEDURAL_HISTORY = "procedural_history"
    ISSUE = "issue"
    HOLDING = "holding"
    REASONING = "reasoning"
    RULE = "rule"
    DISSENT = "dissent"
    CONCURRENCE = "concurrence"
    UNKNOWN = "unknown"


# ============================================================================
# SECTION DETECTION PATTERNS
# ============================================================================

# Compiled patterns for detecting different sections
SECTION_PATTERNS = {
    SectionType.FACTS: re.compile(
        r'\b(?:FACTS?|BACKGROUND|FACTUAL BACKGROUND|'
        r'The facts? (?:are|of this case)|'
        r'This case arises from|'
        r'This appeal arises from)\b',
        re.IGNORECASE
    ),
    SectionType.PROCEDURAL_HISTORY: re.compile(
        r'\b(?:PROCEDURAL (?:HISTORY|BACKGROUND|POSTURE)|'
        r'The (?:district|trial) court|'
        r'(?:Plaintiff|Defendant) (?:filed|moved|appealed)|'
        r'The court below|'
        r'The lower court)\b',
        re.IGNORECASE
    ),
    SectionType.ISSUE: re.compile(
        r'\b(?:ISSUE|QUESTION PRESENTED|'
        r'The (?:sole |)(?:issue|question) (?:presented |)(?:is|before)|'
        r'We (?:must|are asked to) (?:decide|determine)|'
        r'Whether)\b',
        re.IGNORECASE
    ),
    SectionType.HOLDING: re.compile(
        r'\b(?:HOLDING|CONCLUSION|'
        r'We hold that|'
        r'We conclude that|'
        r'We find that|'
        r'We (?:affirm|reverse|remand)|'
        r'(?:Affirmed|Reversed|Remanded)\.?$)\b',
        re.IGNORECASE | re.MULTILINE
    ),
    SectionType.REASONING: re.compile(
        r'\b(?:ANALYSIS|DISCUSSION|REASONING|'
        r'We (?:begin|turn|look) (?:by|to|at)|'
        r'The question (?:thus |)(?:is|becomes)|'
        r'(?:First|Second|Third|Finally),)\b',
        re.IGNORECASE
    ),
    SectionType.RULE: re.compile(
        r'\b(?:Under (?:the |this )?(?:rule|standard|test)|'
        r'The (?:applicable |relevant )?(?:rule|standard|test) (?:is|requires)|'
        r'To establish|'
        r'must (?:prove|show|demonstrate))\b',
        re.IGNORECASE
    ),
    SectionType.DISSENT: re.compile(
        r'\b(?:DISSENT|DISSENTING OPINION|'
        r'(?:Justice|Judge) .+ dissenting|'
        r', dissenting:?)\b',
        re.IGNORECASE
    ),
    SectionType.CONCURRENCE: re.compile(
        r'\b(?:CONCURRENCE|CONCURRING OPINION|'
        r'(?:Justice|Judge) .+ concurring|'
        r', concurring:?)\b',
        re.IGNORECASE
    ),
}


# Signal words for different components
ISSUE_SIGNALS = [
    "whether", "issue", "question", "we must decide",
    "court must determine", "question presented"
]

HOLDING_SIGNALS = [
    "we hold", "we conclude", "we find", "we reverse",
    "we affirm", "we remand", "held that", "concluded that"
]

RULE_SIGNALS = [
    "the rule", "the test", "the standard", "to establish",
    "requires that", "must prove", "elements are"
]

REASONING_SIGNALS = [
    "because", "therefore", "thus", "accordingly",
    "first", "second", "third", "moreover", "furthermore",
    "in addition", "however", "nevertheless"
]


# ============================================================================
# CASE BRIEF EXTRACTION
# ============================================================================

def extract_case_brief(text: str) -> Dict[str, any]:
    """
    Extract a case brief from opinion text.

    This is a pattern-based extraction. For best results with AI-enhanced
    extraction, this should be used with Claude via the MCP server.

    Returns:
        Dictionary with brief components
    """
    # Extract basic metadata
    citation = extract_citation_from_header(text)
    case_name = extract_case_name(text)
    court_year = extract_court_and_year(text)

    # Detect sections
    sections = detect_sections(text)

    # Extract components
    facts = extract_facts(text, sections)
    procedural_history = extract_procedural_history(text, sections)
    issues = extract_issues(text, sections)
    holdings = extract_holdings(text, sections)
    reasoning = extract_reasoning(text, sections)
    rules = extract_rules(text, sections)
    dissents = extract_dissents(text, sections)

    return {
        "case_name": case_name,
        "citation": citation,
        "court": court_year.get("court", ""),
        "year": court_year.get("year", ""),
        "facts": facts,
        "procedural_history": procedural_history,
        "issues": issues,
        "holdings": holdings,
        "reasoning": reasoning,
        "rules": rules,
        "dissents": dissents,
        "sections_detected": list(sections.keys())
    }


def detect_sections(text: str) -> Dict[SectionType, List[Tuple[int, int]]]:
    """
    Detect different sections in the opinion text.

    Returns:
        Dictionary mapping section types to list of (start, end) positions
    """
    sections = {}

    for section_type, pattern in SECTION_PATTERNS.items():
        matches = []
        for match in pattern.finditer(text):
            matches.append((match.start(), match.end()))

        if matches:
            sections[section_type] = matches

    return sections


def extract_facts(text: str, sections: Dict) -> List[str]:
    """Extract factual statements from the opinion"""
    facts = []

    # Look for facts section
    if SectionType.FACTS in sections:
        for start, _ in sections[SectionType.FACTS]:
            # Get next 500 characters as facts section
            end = min(start + 1000, len(text))
            facts_section = text[start:end]

            # Split into sentences
            sentences = split_into_sentences(facts_section)
            facts.extend(sentences[:5])  # Take first 5 sentences

    # Also look for fact patterns
    fact_patterns = [
        r'(?:The|On) (?:\w+ \d+, \d{4}|[A-Z][a-z]+ \d+, \d{4})[,\.]',  # Dates
        r'(?:Plaintiff|Defendant|Appellant|Appellee) \w+',  # Party actions
    ]

    for pattern in fact_patterns:
        matches = re.finditer(pattern, text)
        for match in matches:
            # Get surrounding sentence
            sent_start = text.rfind('.', 0, match.start()) + 1
            sent_end = text.find('.', match.end()) + 1
            if sent_end > sent_start:
                sentence = text[sent_start:sent_end].strip()
                if sentence and len(sentence) < 300:
                    facts.append(sentence)

    # Remove duplicates while preserving order
    seen = set()
    unique_facts = []
    for fact in facts:
        if fact not in seen:
            seen.add(fact)
            unique_facts.append(fact)

    return unique_facts[:10]  # Return top 10


def extract_procedural_history(text: str, sections: Dict) -> List[str]:
    """Extract procedural history"""
    history = []

    # Look for procedural history section
    if SectionType.PROCEDURAL_HISTORY in sections:
        for start, _ in sections[SectionType.PROCEDURAL_HISTORY]:
            end = min(start + 800, len(text))
            proc_section = text[start:end]

            sentences = split_into_sentences(proc_section)
            history.extend(sentences[:4])

    # Look for procedural keywords
    proc_keywords = [
        r'(?:district|trial) court (?:granted|denied|dismissed)',
        r'(?:Plaintiff|Defendant) (?:filed|moved)',
        r'court (?:entered|issued)',
        r'(?:granted|denied) (?:summary judgment|motion to dismiss)',
    ]

    for pattern in proc_keywords:
        matches = re.finditer(pattern, text, re.IGNORECASE)
        for match in matches:
            sentence = extract_sentence_containing(text, match.start())
            if sentence:
                history.append(sentence)

    # Remove duplicates
    return list(dict.fromkeys(history))[:8]


def extract_issues(text: str, sections: Dict) -> List[str]:
    """Extract the legal issues presented"""
    issues = []

    # Look for issue section
    if SectionType.ISSUE in sections:
        for start, _ in sections[SectionType.ISSUE]:
            end = min(start + 500, len(text))
            issue_section = text[start:end]

            sentences = split_into_sentences(issue_section)
            issues.extend(sentences[:3])

    # Look for "whether" questions
    whether_pattern = re.compile(r'\b[Ww]hether\b[^.?]+[.?]')
    for match in whether_pattern.finditer(text):
        issues.append(match.group(0).strip())

    # Look for "issue is" statements
    issue_is_pattern = re.compile(r'[Tt]he (?:sole |dispositive |)(?:issue|question)[^.]+\.', re.IGNORECASE)
    for match in issue_is_pattern.finditer(text):
        issues.append(match.group(0).strip())

    return list(dict.fromkeys(issues))[:5]


def extract_holdings(text: str, sections: Dict) -> List[str]:
    """Extract the court's holdings"""
    holdings = []

    # Look for holding section
    if SectionType.HOLDING in sections:
        for start, _ in sections[SectionType.HOLDING]:
            end = min(start + 400, len(text))
            holding_section = text[start:end]

            sentences = split_into_sentences(holding_section)
            holdings.extend(sentences[:2])

    # Look for holding signals
    for signal in HOLDING_SIGNALS:
        pattern = re.compile(r'\b' + re.escape(signal) + r'\b[^.]+\.', re.IGNORECASE)
        for match in pattern.finditer(text):
            holdings.append(match.group(0).strip())

    return list(dict.fromkeys(holdings))[:5]


def extract_reasoning(text: str, sections: Dict) -> List[str]:
    """Extract key reasoning points"""
    reasoning = []

    # Look for reasoning section
    if SectionType.REASONING in sections:
        for start, _ in sections[SectionType.REASONING]:
            end = min(start + 1000, len(text))
            reasoning_section = text[start:end]

            sentences = split_into_sentences(reasoning_section)
            reasoning.extend(sentences[:6])

    # Look for reasoning signals
    for signal in REASONING_SIGNALS:
        pattern = re.compile(r'\b' + re.escape(signal) + r'\b[^.]+\.', re.IGNORECASE)
        matches = list(pattern.finditer(text))
        for match in matches[:3]:  # Limit per signal
            reasoning.append(match.group(0).strip())

    return list(dict.fromkeys(reasoning))[:10]


def extract_rules(text: str, sections: Dict) -> List[str]:
    """Extract legal rules and standards"""
    rules = []

    # Look for rule section
    if SectionType.RULE in sections:
        for start, _ in sections[SectionType.RULE]:
            end = min(start + 500, len(text))
            rule_section = text[start:end]

            sentences = split_into_sentences(rule_section)
            rules.extend(sentences[:3])

    # Look for rule signals
    for signal in RULE_SIGNALS:
        pattern = re.compile(r'\b' + re.escape(signal) + r'\b[^.]+\.', re.IGNORECASE)
        for match in pattern.finditer(text):
            rules.append(match.group(0).strip())

    # Look for element lists
    element_pattern = re.compile(
        r'(?:elements|factors|requirements|prongs) (?:are|of)[^.]+\(\d+\)[^.]+\.',
        re.IGNORECASE
    )
    for match in element_pattern.finditer(text):
        rules.append(match.group(0).strip())

    return list(dict.fromkeys(rules))[:8]


def extract_dissents(text: str, sections: Dict) -> List[str]:
    """Extract dissenting opinions"""
    dissents = []

    if SectionType.DISSENT in sections:
        for start, _ in sections[SectionType.DISSENT]:
            end = min(start + 600, len(text))
            dissent_section = text[start:end]

            sentences = split_into_sentences(dissent_section)
            dissents.extend(sentences[:4])

    return dissents


# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

def split_into_sentences(text: str) -> List[str]:
    """Split text into sentences"""
    # Simple sentence splitter (can be enhanced)
    sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', text)
    return [s.strip() for s in sentences if len(s.strip()) > 20]


def extract_sentence_containing(text: str, position: int) -> Optional[str]:
    """Extract the full sentence containing the given position"""
    # Find sentence boundaries
    sent_start = text.rfind('.', 0, position) + 1
    sent_end = text.find('.', position) + 1

    if sent_end > sent_start and sent_end > 0:
        return text[sent_start:sent_end].strip()

    return None


def extract_citation_from_header(text: str) -> str:
    """Extract the citation from the header of an opinion"""
    # Look for pattern like "123 F.3d 456"
    citation_pattern = re.compile(r'\d{1,3}\s+[A-Z][A-Za-z0-9.\s]+\d+')
    match = citation_pattern.search(text[:500])

    return match.group(0) if match else ""


def extract_case_name(text: str) -> str:
    """Extract case name from opinion header"""
    # Look for "X v. Y" pattern in first 300 chars
    case_pattern = re.compile(r'([A-Z][A-Za-z\s&,.]+)\s+v\.?\s+([A-Z][A-Za-z\s&,.]+)')
    match = case_pattern.search(text[:300])

    if match:
        return f"{match.group(1).strip()} v. {match.group(2).strip()}"

    return ""


def extract_court_and_year(text: str) -> Dict[str, str]:
    """Extract court and year from opinion header"""
    result = {"court": "", "year": ""}

    # Look for year in parentheses
    year_pattern = re.compile(r'\((\d{4})\)')
    year_match = year_pattern.search(text[:500])
    if year_match:
        result["year"] = year_match.group(1)

    # Look for court abbreviations
    court_patterns = [
        (r'\bU\.?S\.?\b', 'U.S. Supreme Court'),
        (r'\b(\d+)(?:st|nd|rd|th)\s+Cir\.?\b', r'\1 Circuit'),
        (r'\bD\.C\.\s+Cir\.?\b', 'D.C. Circuit'),
        (r'\bFed\.?\s+Cir\.?\b', 'Federal Circuit'),
    ]

    for pattern, replacement in court_patterns:
        match = re.search(pattern, text[:500])
        if match:
            if callable(replacement):
                result["court"] = match.expand(replacement)
            else:
                result["court"] = replacement
            break

    return result


# ============================================================================
# ARGUMENT OUTLINE BUILDER
# ============================================================================

@dataclass
class ArgumentPoint:
    """A single point in a legal argument"""
    level: int  # 1 = main point, 2 = subpoint, etc.
    text: str
    supporting_citations: List[str]
    page_number: Optional[int]


def extract_argument_outline(text: str, citations: Dict[str, Dict]) -> List[ArgumentPoint]:
    """
    Extract the argument structure from a brief or opinion.

    Returns:
        List of ArgumentPoint objects representing the argument outline
    """
    outline = []

    # Look for numbered/lettered sections (I., A., 1., a., etc.)
    section_patterns = [
        (1, re.compile(r'^[IVX]+\.', re.MULTILINE)),  # Roman numerals
        (2, re.compile(r'^[A-Z]\.', re.MULTILINE)),  # Capital letters
        (3, re.compile(r'^\d+\.', re.MULTILINE)),  # Numbers
        (4, re.compile(r'^[a-z]\.', re.MULTILINE)),  # Lowercase letters
    ]

    lines = text.split('\n')

    for i, line in enumerate(lines):
        line = line.strip()

        # Check against each pattern
        for level, pattern in section_patterns:
            if pattern.match(line):
                # This is a section header
                # Get the full text (may span multiple lines)
                full_text = line

                # Look ahead for continuation
                j = i + 1
                while j < len(lines) and lines[j].strip() and not any(
                    p[1].match(lines[j].strip()) for p in section_patterns
                ):
                    full_text += " " + lines[j].strip()
                    j += 1

                # Find citations in this section
                section_citations = []
                for cite in citations.keys():
                    if cite in full_text:
                        section_citations.append(cite)

                outline.append(ArgumentPoint(
                    level=level,
                    text=full_text,
                    supporting_citations=section_citations,
                    page_number=None  # Would need page tracking
                ))

    return outline


def format_argument_outline(outline: List[ArgumentPoint]) -> str:
    """Format an argument outline for display"""
    output = []
    output.append("ARGUMENT OUTLINE")
    output.append("=" * 80)
    output.append("")

    for point in outline:
        indent = "  " * (point.level - 1)
        output.append(f"{indent}{point.text}")

        if point.supporting_citations:
            cite_indent = "  " * point.level
            for cite in point.supporting_citations:
                output.append(f"{cite_indent}  → {cite}")

        output.append("")

    return "\n".join(output)
