"""
Enhanced Citation Extraction with Relationship Mapping
=======================================================
Builds a citation knowledge graph with:
- Relationship tracking (followed, distinguished, overruled, cited)
- Authority strength indicators (controlling, persuasive, dicta)
- Context classification (holding, procedural, background)
- Short form tracking (Id., supra, case name references)
"""

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


class CitationRelationship(Enum):
    """How one case relates to another."""
    FOLLOWED = "followed"
    DISTINGUISHED = "distinguished"
    OVERRULED = "overruled"
    CRITICIZED = "criticized"
    CITED = "cited"  # neutral citation
    ABROGATED = "abrogated"
    SUPERSEDED = "superseded"
    QUESTIONED = "questioned"


class AuthorityStrength(Enum):
    """Strength/weight of the citation."""
    CONTROLLING = "controlling"      # Binding precedent
    HIGHLY_PERSUASIVE = "highly_persuasive"  # Higher court in different jx
    PERSUASIVE = "persuasive"        # Same level court, different jx
    ANALOGOUS = "analogous"          # Similar facts, different area
    BACKGROUND = "background"        # Context only
    DICTA = "dicta"                  # Not part of holding


class CitationContext(Enum):
    """What part of the opinion the citation appears in."""
    HOLDING = "holding"
    RULE_STATEMENT = "rule_statement"
    APPLICATION = "application"
    PROCEDURAL = "procedural"
    BACKGROUND = "background"
    DISSENT = "dissent"
    CONCURRENCE = "concurrence"
    PARENTHETICAL = "parenthetical"


class CitationForm(Enum):
    """Form of the citation."""
    FULL = "full"            # Complete citation
    SHORT = "short"          # Case name only
    ID = "id"                # Id. or Id. at X
    SUPRA = "supra"          # Case name, supra
    HEREINAFTER = "hereinafter"


@dataclass
class EnhancedCitation:
    """Rich citation object with all metadata."""
    full_citation: str
    case_name: Optional[str] = None
    volume: Optional[str] = None
    reporter: Optional[str] = None
    first_page: Optional[str] = None
    pinpoint_cites: List[str] = field(default_factory=list)
    year: Optional[str] = None
    court: Optional[str] = None

    # Relationship data
    relationship: CitationRelationship = CitationRelationship.CITED
    relationship_signal: Optional[str] = None  # The actual text that indicated relationship

    # Strength/weight
    strength: AuthorityStrength = AuthorityStrength.PERSUASIVE
    strength_indicators: List[str] = field(default_factory=list)

    # Context
    context: CitationContext = CitationContext.APPLICATION
    surrounding_text: str = ""

    # Citation form tracking
    form: CitationForm = CitationForm.FULL
    all_references: List[Tuple[int, str]] = field(default_factory=list)  # (position, form)

    # Cross-references
    cites_to: List[str] = field(default_factory=list)  # Cases this case cites
    cited_by: List[str] = field(default_factory=list)  # Cases citing this case

    def to_dict(self) -> Dict:
        """Convert to dictionary for JSON serialization."""
        return {
            'full_citation': self.full_citation,
            'case_name': self.case_name,
            'volume': self.volume,
            'reporter': self.reporter,
            'first_page': self.first_page,
            'pinpoint_cites': self.pinpoint_cites,
            'year': self.year,
            'court': self.court,
            'relationship': self.relationship.value,
            'relationship_signal': self.relationship_signal,
            'strength': self.strength.value,
            'strength_indicators': self.strength_indicators,
            'context': self.context.value,
            'surrounding_text': self.surrounding_text[:200] if self.surrounding_text else "",
            'form': self.form.value,
            'reference_count': len(self.all_references),
            'cites_to': self.cites_to,
            'cited_by': self.cited_by
        }


# ============================================================================
# RELATIONSHIP DETECTION PATTERNS
# ============================================================================

RELATIONSHIP_PATTERNS = {
    CitationRelationship.FOLLOWED: [
        r'(?:following|accord|in accord(?:ance)? with|consistent with|as (?:held|stated) in|reaffirm(?:ed|ing))',
        r'(?:we (?:follow|adopt|apply)|the court (?:followed|adopted|applied))',
        r'(?:binding (?:precedent|authority)|controlling)',
    ],
    CitationRelationship.DISTINGUISHED: [
        r'(?:distinguish(?:ed|ing|able)?|unlike|different from|in contrast to)',
        r'(?:but see|cf\.|contra)',
        r'(?:not applicable|does not apply|inapposite)',
        r'(?:factually distinct|distinguishable on)',
    ],
    CitationRelationship.OVERRULED: [
        r'(?:overrul(?:ed|ing)|no longer good law|abrogated)',
        r'(?:expressly reject(?:ed|ing)|repudiat(?:ed|ing))',
        r'(?:supersed(?:ed|ing) by)',
    ],
    CitationRelationship.CRITICIZED: [
        r'(?:criticiz(?:ed|ing)|question(?:ed|ing)|doubt(?:ed|ing))',
        r'(?:but see|see contra)',
        r'(?:we (?:decline|refuse) to (?:follow|adopt))',
    ],
    CitationRelationship.ABROGATED: [
        r'(?:abrogat(?:ed|ing)|supersed(?:ed|ing))',
        r'(?:no longer (?:valid|applicable|controlling))',
    ],
    CitationRelationship.QUESTIONED: [
        r'(?:question(?:ed|able|ing)|cast(?:ing)? doubt)',
        r'(?:tension with|inconsistent with)',
    ],
}

# Compile patterns
COMPILED_RELATIONSHIP_PATTERNS = {
    rel: [re.compile(p, re.IGNORECASE) for p in patterns]
    for rel, patterns in RELATIONSHIP_PATTERNS.items()
}


# ============================================================================
# STRENGTH INDICATOR PATTERNS
# ============================================================================

STRENGTH_PATTERNS = {
    AuthorityStrength.CONTROLLING: [
        r'(?:controlling|binding|mandatory|dispositive)',
        r'(?:we are bound|must follow|compels)',
        r'(?:directly on point|squarely addresses)',
        r'(?:this court held|we held|our precedent)',
    ],
    AuthorityStrength.HIGHLY_PERSUASIVE: [
        r'(?:highly persuasive|particularly instructive|leading case)',
        r'(?:Supreme Court|Court of Appeals)',
        r'(?:well-established|well-settled|hornbook)',
    ],
    AuthorityStrength.PERSUASIVE: [
        r'(?:persuasive|instructive|helpful|useful)',
        r'(?:other (?:courts|jurisdictions) have)',
        r'(?:we find .+ reasoning persuasive)',
    ],
    AuthorityStrength.ANALOGOUS: [
        r'(?:analog(?:ous|y)|similar(?:ly)?|like)',
        r'(?:by analogy|applying .+ by analogy)',
    ],
    AuthorityStrength.DICTA: [
        r'(?:dicta|dictum|in passing|by the way)',
        r'(?:not necessary to|did not need to decide)',
        r'(?:observation|noting that|suggested)',
    ],
}

COMPILED_STRENGTH_PATTERNS = {
    strength: [re.compile(p, re.IGNORECASE) for p in patterns]
    for strength, patterns in STRENGTH_PATTERNS.items()
}


# ============================================================================
# CONTEXT DETECTION PATTERNS
# ============================================================================

CONTEXT_PATTERNS = {
    CitationContext.HOLDING: [
        r'(?:we hold|the court held|holding that|held that)',
        r'(?:we conclude|concluded that|we decide)',
        r'(?:accordingly|therefore|thus|for (?:these|the foregoing) reasons)',
    ],
    CitationContext.RULE_STATEMENT: [
        r'(?:the (?:rule|test|standard) is|under .+ a (?:party|plaintiff|defendant) must)',
        r'(?:requires that|must (?:show|prove|demonstrate|establish))',
        r'(?:elements (?:are|include)|factors (?:are|include))',
    ],
    CitationContext.PROCEDURAL: [
        r'(?:procedural(?:ly)?|motion (?:to|for)|appeal(?:ed)?|affirm(?:ed)?|revers(?:ed|ing))',
        r'(?:jurisdiction|venue|standing|mootness|ripeness)',
        r'(?:remand(?:ed)?|dismiss(?:ed)?|summary judgment)',
    ],
    CitationContext.DISSENT: [
        r'(?:dissent(?:ing|s)?|I (?:would|disagree)|the majority)',
    ],
    CitationContext.CONCURRENCE: [
        r'(?:concur(?:ring|rence)?|I (?:agree|join)|separately)',
    ],
}

COMPILED_CONTEXT_PATTERNS = {
    ctx: [re.compile(p, re.IGNORECASE) for p in patterns]
    for ctx, patterns in CONTEXT_PATTERNS.items()
}


# ============================================================================
# CITATION PATTERNS (Enhanced)
# ============================================================================

# Case name pattern - captures party names before citation
CASE_NAME_PATTERN = re.compile(
    r'([A-Z][a-zA-Z\'\-]+(?:\s+(?:v\.?|vs\.?)\s+[A-Z][a-zA-Z\'\-]+(?:\s+(?:of|&|and|the|Inc\.|Corp\.|Co\.|LLC|Ltd\.?|et\s+al\.?))*)?),?\s*'
    r'(\d{1,3})\s+'
    r'([A-Z][a-zA-Z\.]+(?:\s+(?:2d|3d|4th|App\.?|Ct\.?|Supp\.?))*)\s+'
    r'(\d+)'
    r'(?:,?\s*(\d+))?'  # pinpoint
    r'(?:\s*\(([^)]+)\))?',  # parenthetical with year/court
    re.MULTILINE
)

# Full citation without case name
FULL_CITE_PATTERN = re.compile(
    r'(\d{1,3})\s+'
    r'([A-Z][a-zA-Z\.]+(?:\s+(?:2d|3d|4th|App\.?|Ct\.?|Supp\.?))*)\s+'
    r'(\d+)'
    r'(?:,?\s*(?:at\s+)?(\d+(?:\s*[-–]\s*\d+)?))?'
    r'(?:\s*\(([^)]+)\))?'
)

# Id. citations
ID_CITE_PATTERN = re.compile(r'\b[Ii]d\.(?:\s+at\s+(\d+(?:\s*[-–]\s*\d+)?))?')

# Supra citations
SUPRA_PATTERN = re.compile(r'([A-Z][a-zA-Z\'\-]+),?\s+supra(?:,?\s+at\s+(\d+))?', re.IGNORECASE)

# Short form (case name only)
SHORT_FORM_PATTERN = re.compile(r'([A-Z][a-zA-Z\'\-]+),?\s+(\d+)\s+([A-Z][a-zA-Z\.]+)\s+at\s+(\d+)')


# ============================================================================
# MAIN EXTRACTION FUNCTION
# ============================================================================

def extract_enhanced_citations(text: str) -> Dict[str, EnhancedCitation]:
    """
    Extract citations with full relationship and strength analysis.

    Args:
        text: Document text to analyze

    Returns:
        Dictionary mapping citation key -> EnhancedCitation object
    """
    citations: Dict[str, EnhancedCitation] = {}
    citation_positions: List[Tuple[int, str, str]] = []  # (position, key, form)
    last_citation_key = None

    # Clean text but preserve structure
    clean_text = text.replace('\r\n', '\n').replace('\r', '\n')

    # ========================================================================
    # PASS 1: Find all full citations
    # ========================================================================

    for match in FULL_CITE_PATTERN.finditer(clean_text):
        volume = match.group(1)
        reporter = match.group(2)
        first_page = match.group(3)
        pinpoint = match.group(4)
        paren = match.group(5)

        # Build citation key
        cite_key = f"{volume} {reporter} {first_page}"

        # Get surrounding context (200 chars before and after)
        start_pos = max(0, match.start() - 200)
        end_pos = min(len(clean_text), match.end() + 200)
        context = clean_text[start_pos:end_pos]

        # Try to find case name before citation
        case_name = None
        pre_text = clean_text[max(0, match.start() - 150):match.start()]
        # Pattern for "Party v. Party" or "Party Name v. Other Party Name"
        # Handles multi-word party names, common suffixes (Inc., Corp., etc.)
        name_match = re.search(
            r'([A-Z][a-zA-Z\'\-\.]+(?:\s+(?:of|&|and|the|Inc\.|Corp\.|Co\.|LLC|Ltd\.|ex\s+rel\.?|in\s+re|et\s+al\.?|[A-Z][a-z]+))*'
            r'(?:\s+v\.?\s+'
            r'[A-Z][a-zA-Z\'\-\.]+(?:\s+(?:of|&|and|the|Inc\.|Corp\.|Co\.|LLC|Ltd\.|et\s+al\.?|[A-Z][a-z]+))*)?)'
            r',?\s*$',
            pre_text
        )
        if name_match:
            case_name = name_match.group(1).strip()
            # Clean up trailing punctuation
            case_name = re.sub(r'[,\s]+$', '', case_name)

        # Parse year and court from parenthetical
        year = None
        court = None
        if paren:
            year_match = re.search(r'((?:19|20)\d{2})', paren)
            if year_match:
                year = year_match.group(1)
            court = paren.replace(year or '', '').strip().strip('()')

        # Create or update citation
        if cite_key not in citations:
            citations[cite_key] = EnhancedCitation(
                full_citation=match.group(0),
                case_name=case_name,
                volume=volume,
                reporter=reporter,
                first_page=first_page,
                year=year,
                court=court,
                surrounding_text=context,
                form=CitationForm.FULL
            )

        cite = citations[cite_key]

        # Add pinpoint
        if pinpoint and pinpoint not in cite.pinpoint_cites:
            cite.pinpoint_cites.append(pinpoint)

        # Track position
        cite.all_references.append((match.start(), 'full'))
        citation_positions.append((match.start(), cite_key, 'full'))
        last_citation_key = cite_key

        # Detect relationship from context
        pre_context = clean_text[max(0, match.start() - 100):match.start()]
        relationship, signal = detect_relationship(pre_context)
        if relationship != CitationRelationship.CITED:
            cite.relationship = relationship
            cite.relationship_signal = signal

        # Detect strength
        strength, indicators = detect_strength(context)
        if strength != AuthorityStrength.PERSUASIVE:
            cite.strength = strength
            cite.strength_indicators = indicators

        # Detect context
        cite.context = detect_context(context)

    # ========================================================================
    # PASS 2: Find Id. citations and link to last cited case
    # ========================================================================

    for match in ID_CITE_PATTERN.finditer(clean_text):
        if last_citation_key and last_citation_key in citations:
            cite = citations[last_citation_key]
            cite.all_references.append((match.start(), 'id'))
            citation_positions.append((match.start(), last_citation_key, 'id'))

            # Add pinpoint from Id. at X
            pinpoint = match.group(1)
            if pinpoint and pinpoint not in cite.pinpoint_cites:
                cite.pinpoint_cites.append(pinpoint)

    # ========================================================================
    # PASS 3: Find supra citations
    # ========================================================================

    for match in SUPRA_PATTERN.finditer(clean_text):
        case_fragment = match.group(1)
        pinpoint = match.group(2)

        # Try to match to existing citation by case name
        for key, cite in citations.items():
            if cite.case_name and case_fragment.lower() in cite.case_name.lower():
                cite.all_references.append((match.start(), 'supra'))
                citation_positions.append((match.start(), key, 'supra'))
                if pinpoint and pinpoint not in cite.pinpoint_cites:
                    cite.pinpoint_cites.append(pinpoint)
                break

    # ========================================================================
    # PASS 4: Build citation network (which cases cite which)
    # ========================================================================

    # Sort positions to track citation order
    citation_positions.sort(key=lambda x: x[0])

    # For each citation, look for other citations in its context
    for key, cite in citations.items():
        context = cite.surrounding_text
        for other_key, other_cite in citations.items():
            if other_key != key:
                # Check if other citation appears in this citation's context
                if other_cite.volume and other_cite.reporter:
                    other_pattern = f"{other_cite.volume}.*?{other_cite.reporter}"
                    if re.search(other_pattern, context):
                        if other_key not in cite.cites_to:
                            cite.cites_to.append(other_key)
                        if key not in other_cite.cited_by:
                            other_cite.cited_by.append(key)

    return citations


def detect_relationship(pre_context: str) -> Tuple[CitationRelationship, Optional[str]]:
    """
    Detect the relationship signal before a citation.

    Args:
        pre_context: Text immediately before the citation

    Returns:
        Tuple of (relationship type, signal text that indicated it)
    """
    for rel, patterns in COMPILED_RELATIONSHIP_PATTERNS.items():
        for pattern in patterns:
            match = pattern.search(pre_context)
            if match:
                return rel, match.group(0)

    return CitationRelationship.CITED, None


def detect_strength(context: str) -> Tuple[AuthorityStrength, List[str]]:
    """
    Detect authority strength indicators in citation context.

    Args:
        context: Text surrounding the citation

    Returns:
        Tuple of (strength level, list of indicator phrases found)
    """
    indicators = []
    detected_strength = AuthorityStrength.PERSUASIVE

    # Check patterns in order of strength
    for strength in [AuthorityStrength.CONTROLLING, AuthorityStrength.HIGHLY_PERSUASIVE,
                     AuthorityStrength.DICTA, AuthorityStrength.ANALOGOUS]:
        for pattern in COMPILED_STRENGTH_PATTERNS.get(strength, []):
            match = pattern.search(context)
            if match:
                indicators.append(match.group(0))
                if detected_strength == AuthorityStrength.PERSUASIVE:
                    detected_strength = strength

    return detected_strength, indicators


def detect_context(context: str) -> CitationContext:
    """
    Detect what part of the opinion the citation appears in.

    Args:
        context: Text surrounding the citation

    Returns:
        Citation context type
    """
    for ctx, patterns in COMPILED_CONTEXT_PATTERNS.items():
        for pattern in patterns:
            if pattern.search(context):
                return ctx

    return CitationContext.APPLICATION


# ============================================================================
# OUTPUT FORMATTING
# ============================================================================

def format_enhanced_citations_output(citations: Dict[str, EnhancedCitation]) -> str:
    """Format enhanced citations for human-readable output."""
    if not citations:
        return "No citations found."

    output = []
    output.append("ENHANCED CITATION ANALYSIS")
    output.append("=" * 70)
    output.append(f"\nTotal unique citations: {len(citations)}")

    # Group by relationship type
    by_relationship: Dict[CitationRelationship, List[EnhancedCitation]] = defaultdict(list)
    for cite in citations.values():
        by_relationship[cite.relationship].append(cite)

    # Output controlling authorities first
    controlling = [c for c in citations.values() if c.strength == AuthorityStrength.CONTROLLING]
    if controlling:
        output.append("\n" + "-" * 70)
        output.append("CONTROLLING AUTHORITIES")
        output.append("-" * 70)
        for cite in controlling:
            output.append(format_single_citation(cite))

    # Then by relationship
    for rel in [CitationRelationship.FOLLOWED, CitationRelationship.DISTINGUISHED,
                CitationRelationship.OVERRULED, CitationRelationship.CRITICIZED,
                CitationRelationship.CITED]:
        rel_cites = [c for c in by_relationship[rel] if c.strength != AuthorityStrength.CONTROLLING]
        if rel_cites:
            output.append("\n" + "-" * 70)
            output.append(f"{rel.value.upper()} CASES")
            output.append("-" * 70)
            for cite in rel_cites:
                output.append(format_single_citation(cite))

    # Citation network summary
    network_citations = [c for c in citations.values() if c.cites_to or c.cited_by]
    if network_citations:
        output.append("\n" + "=" * 70)
        output.append("CITATION NETWORK")
        output.append("=" * 70)
        for cite in network_citations:
            if cite.cites_to or cite.cited_by:
                name = cite.case_name or cite.full_citation[:30]
                output.append(f"\n• {name}")
                if cite.cites_to:
                    output.append(f"  Cites to: {', '.join(cite.cites_to[:5])}")
                if cite.cited_by:
                    output.append(f"  Cited by: {', '.join(cite.cited_by[:5])}")

    return '\n'.join(output)


def format_single_citation(cite: EnhancedCitation) -> str:
    """Format a single enhanced citation."""
    lines = []

    # Main citation
    if cite.case_name:
        lines.append(f"\n• {cite.case_name}")
        lines.append(f"  {cite.full_citation}")
    else:
        lines.append(f"\n• {cite.full_citation}")

    # Year and court
    if cite.year or cite.court:
        meta = []
        if cite.court:
            meta.append(cite.court)
        if cite.year:
            meta.append(cite.year)
        lines.append(f"  ({', '.join(meta)})")

    # Pinpoints
    if cite.pinpoint_cites:
        sorted_pinpoints = sorted(set(cite.pinpoint_cites),
                                  key=lambda x: int(x.split('-')[0].split('–')[0]))
        lines.append(f"  Pinpoint pages: {', '.join(sorted_pinpoints)}")

    # Strength
    if cite.strength != AuthorityStrength.PERSUASIVE:
        lines.append(f"  Authority: {cite.strength.value}")
        if cite.strength_indicators:
            lines.append(f"    Indicators: {'; '.join(cite.strength_indicators[:3])}")

    # Relationship
    if cite.relationship != CitationRelationship.CITED:
        lines.append(f"  Treatment: {cite.relationship.value}")
        if cite.relationship_signal:
            lines.append(f"    Signal: \"{cite.relationship_signal}\"")

    # Context
    lines.append(f"  Context: {cite.context.value}")

    # Reference count
    if len(cite.all_references) > 1:
        lines.append(f"  Referenced: {len(cite.all_references)} times")

    return '\n'.join(lines)


# ============================================================================
# WRAPPER FOR MCP INTEGRATION
# ============================================================================

def analyze_citations_enhanced(text: str) -> str:
    """
    Main entry point for MCP server.

    Args:
        text: Document text to analyze

    Returns:
        Formatted analysis string
    """
    citations = extract_enhanced_citations(text)
    return format_enhanced_citations_output(citations)


def get_citations_as_dict(text: str) -> Dict[str, Dict]:
    """
    Get citations as dictionary for programmatic use.

    Args:
        text: Document text to analyze

    Returns:
        Dictionary mapping citation key -> citation data dict
    """
    citations = extract_enhanced_citations(text)
    return {key: cite.to_dict() for key, cite in citations.items()}
