"""
Legal Analysis and Writing Tools
Advanced features for law students and legal professionals
"""

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


# ============================================================================
# BLUEBOOK CITATION FORMATTING
# ============================================================================

class CitationType(Enum):
    """Types of legal citations"""
    CASE = "case"
    STATUTE = "statute"
    CONSTITUTIONAL = "constitutional"
    REGULATION = "regulation"
    TREATISE = "treatise"
    LAW_REVIEW = "law_review"


@dataclass
class BluebookCitation:
    """Represents a properly formatted Bluebook citation"""
    original: str
    formatted: str
    type: CitationType
    errors: List[str]
    suggestions: List[str]


def format_case_citation_bluebook(citation: str) -> BluebookCitation:
    """
    Format a case citation according to Bluebook rules.

    Rules applied:
    - Case name in italics (represented by underscores)
    - Volume, reporter, first page
    - Pinpoint cite if present
    - Parenthetical with court and year
    """
    errors = []
    suggestions = []

    # Extract components
    # Pattern: Volume Reporter Page, Pinpoint (Court Year)
    case_pattern = re.compile(
        r'(\d+)\s+'  # Volume
        r'([A-Z][A-Za-z0-9.\s]+?)\s+'  # Reporter
        r'(\d+)'  # First page
        r'(?:,\s*(\d+(?:-\d+)?))?\s*'  # Optional pinpoint
        r'(?:\(([^)]+)\))?\s*'  # Optional court/year paren
    )

    match = case_pattern.search(citation)
    if not match:
        return BluebookCitation(
            original=citation,
            formatted=citation,
            type=CitationType.CASE,
            errors=["Could not parse citation format"],
            suggestions=["Ensure format is: Volume Reporter Page (Court Year)"]
        )

    volume, reporter, first_page, pinpoint, court_year = match.groups()

    # Clean up reporter abbreviation
    reporter = reporter.strip()

    # Check for common reporter abbreviation errors
    reporter_corrections = {
        'F.2d': 'F.2d',
        'F.3d': 'F.3d',
        'F.4th': 'F.4th',
        'F. 2d': 'F.2d',
        'F. 3d': 'F.3d',
        'F Supp': 'F. Supp.',
        'F.Supp': 'F. Supp.',
        'F. Supp': 'F. Supp.',
        'U.S.': 'U.S.',
        'US': 'U.S.',
        'S.Ct.': 'S. Ct.',
        'S. Ct': 'S. Ct.',
    }

    corrected_reporter = reporter_corrections.get(reporter, reporter)
    if corrected_reporter != reporter:
        suggestions.append(f"Reporter should be '{corrected_reporter}' not '{reporter}'")
        reporter = corrected_reporter

    # Build formatted citation
    formatted_parts = [volume, reporter, first_page]

    if pinpoint:
        formatted_parts.append(f", {pinpoint}")

    if court_year:
        formatted_parts.append(f" ({court_year})")
    else:
        errors.append("Missing court and year in parenthetical")
        suggestions.append("Add (Court Year) at end of citation")

    formatted = ' '.join(formatted_parts)

    # Check spacing
    if '  ' in formatted:
        errors.append("Extra spaces in citation")
        formatted = ' '.join(formatted.split())

    return BluebookCitation(
        original=citation,
        formatted=formatted,
        type=CitationType.CASE,
        errors=errors,
        suggestions=suggestions
    )


# ============================================================================
# TABLE OF AUTHORITIES GENERATOR
# ============================================================================

@dataclass
class TOAEntry:
    """Entry for Table of Authorities"""
    citation: str
    category: str  # "Cases", "Statutes", "Constitutional Provisions", etc.
    pages: List[int]
    formatted_citation: str

    def __lt__(self, other):
        """Sort alphabetically by formatted citation"""
        return self.formatted_citation.lower() < other.formatted_citation.lower()


def generate_table_of_authorities(
    citations: Dict[str, Dict],
    statutes: Dict[str, Dict],
    document_name: str = ""
) -> str:
    """
    Generate a properly formatted Table of Authorities.

    Args:
        citations: Dictionary of case citations from extract_citations
        statutes: Dictionary of statute citations from extract_statutes
        document_name: Optional name of the document

    Returns:
        Formatted Table of Authorities as a string
    """
    output = []
    output.append("TABLE OF AUTHORITIES")
    output.append("=" * 80)
    output.append("")

    # Sort into categories
    cases = []
    constitutional = []
    statutes_list = []
    regulations = []
    other = []

    # Process case citations
    for cite, data in citations.items():
        pages = sorted(set(data.get('pdf_pages', [])))

        # Format pages as "1, 3, 5-7, 10"
        page_ranges = format_page_ranges(pages)

        cases.append(TOAEntry(
            citation=cite,
            category="Cases",
            pages=pages,
            formatted_citation=cite
        ))

    # Process statutes
    for statute, data in statutes.items():
        pages = sorted(set(data.get('pages', [])))

        # Categorize
        if 'U.S. Const.' in statute or 'Const.' in statute:
            constitutional.append(TOAEntry(
                citation=statute,
                category="Constitutional Provisions",
                pages=pages,
                formatted_citation=statute
            ))
        elif 'C.F.R.' in statute or 'Fed. Reg.' in statute:
            regulations.append(TOAEntry(
                citation=statute,
                category="Regulations",
                pages=pages,
                formatted_citation=statute
            ))
        else:
            statutes_list.append(TOAEntry(
                citation=statute,
                category="Statutes",
                pages=pages,
                formatted_citation=statute
            ))

    # Sort each category
    cases.sort()
    constitutional.sort()
    statutes_list.sort()
    regulations.sort()

    # Generate output for each category
    categories = [
        ("Cases", cases),
        ("Constitutional Provisions", constitutional),
        ("Statutes", statutes_list),
        ("Regulations", regulations),
    ]

    for category_name, entries in categories:
        if not entries:
            continue

        output.append(f"{category_name}")
        output.append("-" * 80)
        output.append("")

        for entry in entries:
            page_str = format_page_ranges(entry.pages)
            # Format with proper indentation and page alignment
            citation_part = entry.formatted_citation

            # Calculate spacing for alignment (aim for page numbers at column 60)
            target_col = 60
            citation_len = len(citation_part)

            if citation_len < target_col - 5:
                spacing = " " * (target_col - citation_len)
                output.append(f"{citation_part}{spacing}{page_str}")
            else:
                # Citation is too long, put pages on next line
                output.append(f"{citation_part}")
                output.append(f"{'': >{target_col}}{page_str}")

        output.append("")

    return "\n".join(output)


def format_page_ranges(pages: List[int]) -> str:
    """
    Format page numbers into ranges.
    Example: [1, 2, 3, 5, 7, 8, 9] -> "1-3, 5, 7-9"
    """
    if not pages:
        return ""

    pages = sorted(set(pages))
    ranges = []
    start = pages[0]
    end = pages[0]

    for i in range(1, len(pages)):
        if pages[i] == end + 1:
            end = pages[i]
        else:
            if start == end:
                ranges.append(str(start))
            elif end == start + 1:
                ranges.append(f"{start}, {end}")
            else:
                ranges.append(f"{start}-{end}")
            start = pages[i]
            end = pages[i]

    # Add the last range
    if start == end:
        ranges.append(str(start))
    elif end == start + 1:
        ranges.append(f"{start}, {end}")
    else:
        ranges.append(f"{start}-{end}")

    return ", ".join(ranges)


# ============================================================================
# JURISDICTION CLASSIFICATION
# ============================================================================

class Jurisdiction(Enum):
    """Federal and state jurisdictions"""
    SUPREME_COURT = "U.S. Supreme Court"
    FIRST_CIRCUIT = "1st Circuit"
    SECOND_CIRCUIT = "2nd Circuit"
    THIRD_CIRCUIT = "3rd Circuit"
    FOURTH_CIRCUIT = "4th Circuit"
    FIFTH_CIRCUIT = "5th Circuit"
    SIXTH_CIRCUIT = "6th Circuit"
    SEVENTH_CIRCUIT = "7th Circuit"
    EIGHTH_CIRCUIT = "8th Circuit"
    NINTH_CIRCUIT = "9th Circuit"
    TENTH_CIRCUIT = "10th Circuit"
    ELEVENTH_CIRCUIT = "11th Circuit"
    DC_CIRCUIT = "D.C. Circuit"
    FEDERAL_CIRCUIT = "Federal Circuit"
    DISTRICT_COURT = "Federal District Court"
    STATE_SUPREME = "State Supreme Court"
    STATE_APPELLATE = "State Appellate Court"
    STATE_TRIAL = "State Trial Court"
    UNKNOWN = "Unknown"


def classify_jurisdiction(citation: str) -> Jurisdiction:
    """
    Determine the jurisdiction/court level from a citation.
    """
    citation = citation.upper()

    # Supreme Court
    if 'U.S.' in citation and ('S. CT.' not in citation and 'S.CT.' not in citation):
        return Jurisdiction.SUPREME_COURT
    if 'S. CT.' in citation or 'S.CT.' in citation:
        return Jurisdiction.SUPREME_COURT

    # Circuit courts
    circuit_patterns = {
        '1ST CIR': Jurisdiction.FIRST_CIRCUIT,
        '2D CIR': Jurisdiction.SECOND_CIRCUIT,
        '2ND CIR': Jurisdiction.SECOND_CIRCUIT,
        '3D CIR': Jurisdiction.THIRD_CIRCUIT,
        '3RD CIR': Jurisdiction.THIRD_CIRCUIT,
        '4TH CIR': Jurisdiction.FOURTH_CIRCUIT,
        '5TH CIR': Jurisdiction.FIFTH_CIRCUIT,
        '6TH CIR': Jurisdiction.SIXTH_CIRCUIT,
        '7TH CIR': Jurisdiction.SEVENTH_CIRCUIT,
        '8TH CIR': Jurisdiction.EIGHTH_CIRCUIT,
        '9TH CIR': Jurisdiction.NINTH_CIRCUIT,
        '10TH CIR': Jurisdiction.TENTH_CIRCUIT,
        '11TH CIR': Jurisdiction.ELEVENTH_CIRCUIT,
        'D.C. CIR': Jurisdiction.DC_CIRCUIT,
        'FED. CIR': Jurisdiction.FEDERAL_CIRCUIT,
    }

    for pattern, jurisdiction in circuit_patterns.items():
        if pattern in citation:
            return jurisdiction

    # Federal district courts
    if 'F. SUPP' in citation or 'F.SUPP' in citation:
        return Jurisdiction.DISTRICT_COURT

    return Jurisdiction.UNKNOWN


def filter_by_jurisdiction(
    citations: Dict[str, Dict],
    jurisdiction: Jurisdiction
) -> Dict[str, Dict]:
    """
    Filter citations to only those from a specific jurisdiction.
    """
    filtered = {}

    for cite, data in citations.items():
        if classify_jurisdiction(cite) == jurisdiction:
            filtered[cite] = data

    return filtered


def group_by_jurisdiction(citations: Dict[str, Dict]) -> Dict[Jurisdiction, List[str]]:
    """
    Group citations by their jurisdiction.

    Returns:
        Dictionary mapping Jurisdiction to list of citation strings
    """
    grouped = defaultdict(list)

    for cite in citations.keys():
        jurisdiction = classify_jurisdiction(cite)
        grouped[jurisdiction].append(cite)

    # Sort within each jurisdiction
    for jurisdiction in grouped:
        grouped[jurisdiction].sort()

    return dict(grouped)


# ============================================================================
# STANDARD OF REVIEW DETECTION
# ============================================================================

class StandardOfReview(Enum):
    """Common standards of review"""
    DE_NOVO = "de novo"
    CLEAR_ERROR = "clear error"
    ABUSE_OF_DISCRETION = "abuse of discretion"
    SUBSTANTIAL_EVIDENCE = "substantial evidence"
    ARBITRARY_CAPRICIOUS = "arbitrary and capricious"
    PLAIN_ERROR = "plain error"
    REASONABLENESS = "reasonableness"
    UNKNOWN = "unknown"


# Compiled patterns for performance
STANDARD_PATTERNS = {
    StandardOfReview.DE_NOVO: re.compile(
        r'\b(?:de novo|plenary|independent)(?:\s+review)?\b',
        re.IGNORECASE
    ),
    StandardOfReview.CLEAR_ERROR: re.compile(
        r'\bclear(?:ly)?\s+erro(?:r|neous)\b',
        re.IGNORECASE
    ),
    StandardOfReview.ABUSE_OF_DISCRETION: re.compile(
        r'\babuse\s+of\s+discretion\b',
        re.IGNORECASE
    ),
    StandardOfReview.SUBSTANTIAL_EVIDENCE: re.compile(
        r'\bsubstantial\s+evidence\b',
        re.IGNORECASE
    ),
    StandardOfReview.ARBITRARY_CAPRICIOUS: re.compile(
        r'\barbitrary\s+(?:and|or|&)\s+capricious\b',
        re.IGNORECASE
    ),
    StandardOfReview.PLAIN_ERROR: re.compile(
        r'\bplain\s+error\b',
        re.IGNORECASE
    ),
    StandardOfReview.REASONABLENESS: re.compile(
        r'\breasonableness(?:\s+review)?\b',
        re.IGNORECASE
    ),
}


def detect_standards_of_review(text: str) -> List[Tuple[StandardOfReview, str, int]]:
    """
    Detect standards of review in text.

    Returns:
        List of tuples: (standard, context_snippet, character_position)
    """
    results = []

    for standard, pattern in STANDARD_PATTERNS.items():
        for match in pattern.finditer(text):
            # Get context (50 chars before and after)
            start = max(0, match.start() - 50)
            end = min(len(text), match.end() + 50)
            context = text[start:end].strip()

            results.append((standard, context, match.start()))

    # Sort by position
    results.sort(key=lambda x: x[2])

    return results


# ============================================================================
# LEGAL ISSUE/TOPIC CLUSTERING
# ============================================================================

# Common legal topics and their keywords
LEGAL_TOPICS = {
    "Civil Procedure": [
        "jurisdiction", "venue", "standing", "joinder", "discovery",
        "summary judgment", "motion to dismiss", "personal jurisdiction",
        "subject matter jurisdiction", "pleading", "Fed. R. Civ. P."
    ],
    "Constitutional Law": [
        "First Amendment", "Second Amendment", "Fourth Amendment", "Fifth Amendment",
        "Fourteenth Amendment", "Due Process", "Equal Protection", "Commerce Clause",
        "Establishment Clause", "Free Exercise", "free speech", "unreasonable search"
    ],
    "Criminal Law": [
        "mens rea", "actus reus", "felony", "misdemeanor", "homicide", "assault",
        "battery", "larceny", "robbery", "burglary", "conspiracy", "accomplice"
    ],
    "Criminal Procedure": [
        "Miranda", "search and seizure", "warrant", "probable cause", "reasonable suspicion",
        "exclusionary rule", "fruit of the poisonous tree", "good faith exception",
        "stop and frisk", "Terry stop"
    ],
    "Contracts": [
        "offer", "acceptance", "consideration", "breach", "specific performance",
        "damages", "rescission", "statute of frauds", "parol evidence", "UCC",
        "meeting of the minds", "mutual assent"
    ],
    "Torts": [
        "negligence", "duty of care", "breach of duty", "causation", "proximate cause",
        "but-for", "strict liability", "intentional tort", "assault", "battery",
        "false imprisonment", "defamation", "res ipsa loquitur"
    ],
    "Property": [
        "real property", "easement", "covenant", "adverse possession", "title",
        "deed", "mortgage", "zoning", "eminent domain", "landlord", "tenant"
    ],
    "Evidence": [
        "hearsay", "relevance", "prejudicial", "authentication", "best evidence",
        "privilege", "attorney-client", "work product", "expert witness", "Fed. R. Evid."
    ],
    "Securities Fraud": [
        "10b-5", "scienter", "material misstatement", "omission", "confidential witness",
        "insider trading", "securities", "fraud on the market", "PSLRA", "pleading"
    ],
    "Administrative Law": [
        "Chevron", "arbitrary and capricious", "notice and comment", "rulemaking",
        "adjudication", "APA", "agency action", "deference"
    ],
}


def cluster_by_topic(text: str) -> Dict[str, int]:
    """
    Analyze text and cluster by legal topics based on keyword frequency.

    Returns:
        Dictionary mapping topic names to keyword match counts
    """
    text_lower = text.lower()
    topic_scores = {}

    for topic, keywords in LEGAL_TOPICS.items():
        score = 0
        for keyword in keywords:
            # Count occurrences of keyword
            score += len(re.findall(r'\b' + re.escape(keyword.lower()) + r'\b', text_lower))

        if score > 0:
            topic_scores[topic] = score

    # Sort by score descending
    return dict(sorted(topic_scores.items(), key=lambda x: x[1], reverse=True))


def identify_primary_topics(text: str, top_n: int = 3) -> List[Tuple[str, int]]:
    """
    Identify the primary legal topics in a document.

    Returns:
        List of (topic, score) tuples for the top N topics
    """
    scores = cluster_by_topic(text)
    sorted_topics = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_topics[:top_n]
