"""
Brief Quality Assurance Tools for Baymax
========================================
Critical tools for catching problems before filing:
1. Record Cite Validator - flags factual claims without record support
2. Negative Treatment Detector - partial Shepardizing without Westlaw
3. Authority Strength Scorer - ranks cases by binding strength
"""

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 RecordCite:
    """A citation to the record."""
    raw_text: str           # Original text: "R. 123", "Doc. 45-2 at 15"
    doc_type: str           # "R", "Doc", "Exh", "Tr", "Dep", etc.
    doc_number: str         # "123", "45-2"
    page: Optional[str]     # Page or line number if present
    context: str            # Surrounding text


@dataclass
class FactualAssertion:
    """A factual claim in the brief that needs record support."""
    text: str               # The assertion
    has_cite: bool          # Whether it has record support
    cite: Optional[RecordCite]  # The cite if present
    sentence_index: int     # Position in document
    severity: str           # "high", "medium", "low" - how important is this fact
    fact_type: str = "material"  # "material", "background", "procedural", "legal"
    suggested_source: Optional[str] = None  # Suggested record doc type


@dataclass
class CiteProblem:
    """A problem with a record citation."""
    cite_text: str
    problem_type: str       # "pleading_cite", "vague_range", "missing_pinpoint", "self_cite"
    severity: str           # "error", "warning", "info"
    explanation: str
    suggestion: str


@dataclass
class RecordValidation:
    """Results of record cite validation."""
    total_assertions: int
    supported_assertions: int
    unsupported_assertions: int
    all_cites: List[RecordCite]
    missing_cite_flags: List[FactualAssertion]
    coverage_percentage: float
    cite_problems: List[CiteProblem] = field(default_factory=list)
    material_facts_unsupported: int = 0
    background_facts_unsupported: int = 0


class TreatmentSignal(Enum):
    """How a later case treats an earlier case."""
    FOLLOWED = "followed"
    APPLIED = "applied"
    CITED = "cited"
    DISTINGUISHED = "distinguished"
    LIMITED = "limited"
    QUESTIONED = "questioned"
    CRITICIZED = "criticized"
    OVERRULED = "overruled"
    SUPERSEDED = "superseded"


@dataclass
class NegativeTreatment:
    """Detected negative treatment of a case."""
    citing_case: str        # Case that cites negatively
    citing_citation: str    # Citation of citing case
    signal: TreatmentSignal
    context_quote: str      # The negative language
    severity: str           # "critical", "serious", "moderate", "minor"
    explanation: str        # What this means


@dataclass
class AuthorityScore:
    """Scored authority for citation ordering."""
    case_name: str
    citation: str
    total_score: int        # 0-100
    court_level_score: int
    jurisdiction_score: int
    recency_score: int
    treatment_score: int
    recommendation: str     # "Lead citation", "Strong support", etc.
    factors: List[str]      # Explanation of scoring


# ============================================================================
# RECORD CITE PATTERNS
# ============================================================================

# Patterns for record citations
RECORD_CITE_PATTERNS = [
    # Standard record cites
    (r'\(R\.\s*(\d+(?:-\d+)?)\)', 'R'),                    # (R. 123) or (R. 123-45)
    (r'R\.\s*(?:at\s*)?(\d+(?:-\d+)?)', 'R'),             # R. 123 or R. at 123
    (r'\(Doc\.\s*(\d+(?:-\d+)?)', 'Doc'),                 # (Doc. 45) or (Doc. 45-2)
    (r'Doc\.\s*(\d+(?:-\d+)?)', 'Doc'),                   # Doc. 45
    (r'Dkt\.\s*(?:No\.\s*)?(\d+)', 'Dkt'),                # Dkt. 123 or Dkt. No. 123

    # Exhibit cites
    (r'\((?:Ex\.|Exh\.|Exhibit)\s*([A-Z0-9]+)\)', 'Exh'), # (Ex. A) or (Exh. 1)
    (r'(?:Ex\.|Exh\.|Exhibit)\s*([A-Z0-9]+)', 'Exh'),     # Ex. A

    # Transcript cites
    (r'\(Tr\.\s*(\d+(?::\d+)?(?:-\d+)?)\)', 'Tr'),        # (Tr. 50:1-10)
    (r'Tr\.\s*(?:at\s*)?(\d+(?::\d+)?)', 'Tr'),           # Tr. 50:1

    # Deposition cites
    (r'\(([A-Z][a-z]+)\s+Dep\.\s*(\d+(?::\d+)?)\)', 'Dep'),  # (Smith Dep. 45:10)
    (r'([A-Z][a-z]+)\s+Dep\.\s*(?:at\s*)?(\d+)', 'Dep'),     # Smith Dep. at 45

    # Appendix/Joint Appendix
    (r'\((?:App\.|JA)\s*(\d+)\)', 'App'),                 # (App. 123) or (JA 123)
    (r'(?:App\.|JA)\s*(?:at\s*)?(\d+)', 'App'),           # App. 123

    # Declaration/Affidavit
    (r'\(([A-Z][a-z]+)\s+(?:Decl\.|Aff\.)\s*(?:¶|para\.)?\s*(\d+)\)', 'Decl'),

    # Complaint/Answer paragraphs
    (r'\(Compl\.\s*(?:¶|para\.)?\s*(\d+)\)', 'Compl'),
    (r'\(Am\.\s*Compl\.\s*(?:¶|para\.)?\s*(\d+)\)', 'Am. Compl'),
    (r'\(Answer\s*(?:¶|para\.)?\s*(\d+)\)', 'Answer'),
]

COMPILED_RECORD_PATTERNS = [(re.compile(p, re.IGNORECASE), t) for p, t in RECORD_CITE_PATTERNS]


# Patterns that suggest factual assertions needing record support
FACTUAL_ASSERTION_INDICATORS = [
    # Temporal facts
    r'\b(?:on|in|during|after|before)\s+(?:January|February|March|April|May|June|July|August|September|October|November|December|\d{1,2}/\d{1,2}/\d{2,4})',
    # Quantitative facts
    r'\b(?:approximately|\$|percent|%|\d+\s+(?:days|months|years|dollars|shares|employees))',
    # Action facts
    r'\b(?:plaintiff|defendant|company|CEO|CFO|he|she|they)\s+(?:stated|said|wrote|sent|received|knew|learned|discovered|admitted|testified|confirmed)',
    # Document facts
    r'\b(?:the\s+)?(?:email|letter|memo|report|filing|document|contract|agreement)\s+(?:stated|showed|revealed|indicated)',
    # Meeting/communication facts
    r'\b(?:during|at|in)\s+(?:the|a)\s+(?:meeting|call|conference|discussion)',
]

COMPILED_FACTUAL_PATTERNS = [re.compile(p, re.IGNORECASE) for p in FACTUAL_ASSERTION_INDICATORS]


# ============================================================================
# NEGATIVE TREATMENT PATTERNS
# ============================================================================

NEGATIVE_TREATMENT_PATTERNS = {
    TreatmentSignal.OVERRULED: [
        r'overruled\s+(?:by|in)',
        r'no\s+longer\s+good\s+law',
        r'expressly\s+overruled',
        r'abrogated\s+by',
    ],
    TreatmentSignal.SUPERSEDED: [
        r'superseded\s+by\s+statute',
        r'superseded\s+by\s+regulation',
        r'legislatively\s+overruled',
    ],
    TreatmentSignal.LIMITED: [
        r'limited\s+(?:to|by)',
        r'narrow(?:ed|ly)\s+(?:the\s+)?holding',
        r'confined\s+to\s+its\s+facts',
        r'decline(?:d|s)?\s+to\s+extend',
        r'(?:did|do)\s+not\s+extend',
    ],
    TreatmentSignal.DISTINGUISHED: [
        r'distinguish(?:ed|able)',
        r'factually\s+distinguishable',
        r'inapposite',
        r'(?:is|are)\s+different\s+from',
        r'unlike\s+(?:in\s+)?[A-Z]',
    ],
    TreatmentSignal.QUESTIONED: [
        r'question(?:ed|able|ing)',
        r'doubt(?:ed|ful)',
        r'cast\s+doubt',
        r'called\s+into\s+question',
        r'criticized\s+(?:by|in)',
    ],
    TreatmentSignal.CRITICIZED: [
        r'criticiz(?:ed|ing)',
        r'disagree(?:d|s|ing)?\s+with',
        r'reject(?:ed|s|ing)',
        r'(?:respectfully\s+)?decline(?:d|s)?\s+to\s+follow',
    ],
}

COMPILED_NEGATIVE_PATTERNS = {
    signal: [re.compile(p, re.IGNORECASE) for p in patterns]
    for signal, patterns in NEGATIVE_TREATMENT_PATTERNS.items()
}

# Severity mapping
TREATMENT_SEVERITY = {
    TreatmentSignal.OVERRULED: "critical",
    TreatmentSignal.SUPERSEDED: "critical",
    TreatmentSignal.LIMITED: "serious",
    TreatmentSignal.QUESTIONED: "moderate",
    TreatmentSignal.CRITICIZED: "moderate",
    TreatmentSignal.DISTINGUISHED: "minor",
}


# ============================================================================
# RECORD CITE EXTRACTION AND VALIDATION
# ============================================================================

def extract_record_cites(text: str) -> List[RecordCite]:
    """
    Extract all record citations from text.

    Args:
        text: Brief or document text

    Returns:
        List of RecordCite objects found
    """
    cites = []
    seen = set()  # Avoid duplicates

    for pattern, doc_type in COMPILED_RECORD_PATTERNS:
        for match in pattern.finditer(text):
            raw = match.group(0)

            # Avoid duplicates
            if raw in seen:
                continue
            seen.add(raw)

            # Extract components
            groups = match.groups()
            doc_number = groups[0] if groups else ""

            # Look for page number after the cite
            page = None
            after_match = text[match.end():match.end()+20]
            page_match = re.search(r'(?:at|p\.|pp\.)\s*(\d+)', after_match)
            if page_match:
                page = page_match.group(1)

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

            cites.append(RecordCite(
                raw_text=raw,
                doc_type=doc_type,
                doc_number=doc_number,
                page=page,
                context=context
            ))

    return cites


def identify_factual_assertions(text: str) -> List[Tuple[str, int]]:
    """
    Identify sentences that appear to be factual assertions needing record support.

    Returns:
        List of (sentence, index) tuples
    """
    assertions = []

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

    for i, sentence in enumerate(sentences):
        # Check if sentence matches factual assertion patterns
        is_factual = any(p.search(sentence) for p in COMPILED_FACTUAL_PATTERNS)

        # Also check for specific factual language
        factual_words = ['stated', 'testified', 'admitted', 'confirmed', 'wrote',
                         'sent', 'received', 'paid', 'earned', 'lost', 'owed']
        has_factual_word = any(word in sentence.lower() for word in factual_words)

        # Skip legal conclusions and argument
        legal_words = ['therefore', 'accordingly', 'thus', 'because', 'since',
                       'the court', 'this court', 'should', 'must', 'requires']
        is_legal = any(word in sentence.lower() for word in legal_words)

        if (is_factual or has_factual_word) and not is_legal and len(sentence) > 30:
            assertions.append((sentence, i))

    return assertions


def validate_record_cites(text: str) -> RecordValidation:
    """
    Validate that factual assertions in brief have record support.

    Enhanced features:
    - Suggests which record doc likely contains missing facts
    - Detects improper citation to pleadings as proof
    - Flags vague citation ranges
    - Distinguishes material facts from background facts

    Args:
        text: Brief text

    Returns:
        RecordValidation with supported/unsupported assertions and problems
    """
    # Extract all record cites
    all_cites = extract_record_cites(text)
    cite_positions = set()
    cite_problems = []

    # Map cite positions
    for pattern, _ in COMPILED_RECORD_PATTERNS:
        for match in pattern.finditer(text):
            # Record the 100 characters before the cite as "covered"
            start = max(0, match.start() - 150)
            for pos in range(start, match.end()):
                cite_positions.add(pos)

    # Check for citation problems
    cite_problems.extend(_detect_pleading_cite_problems(text))
    cite_problems.extend(_detect_vague_citations(text))

    # Find factual assertions
    assertions_raw = identify_factual_assertions(text)
    assertions = []

    # Check each assertion for record support
    for sentence, idx in assertions_raw:
        # Find position in text
        try:
            pos = text.index(sentence)
        except ValueError:
            pos = 0

        # Check if any part of sentence is near a cite
        sentence_end = pos + len(sentence)
        has_cite = any(p in cite_positions for p in range(pos, sentence_end + 50))

        # Classify fact type (material vs. background vs. procedural)
        fact_type = _classify_fact_type(sentence)

        # Determine severity based on content and fact type
        if fact_type == "material":
            if any(word in sentence.lower() for word in ['defendant', 'plaintiff', 'testified', 'admitted', 'knew', 'intended']):
                severity = "high"
            elif any(word in sentence.lower() for word in ['stated', 'wrote', 'sent', 'received']):
                severity = "high"
            else:
                severity = "medium"
        elif fact_type == "background":
            severity = "low"
        else:
            severity = "low"

        # Suggest source document for unsupported facts
        suggested_source = None
        if not has_cite:
            suggested_source = _suggest_record_source(sentence)

        assertions.append(FactualAssertion(
            text=sentence,
            has_cite=has_cite,
            cite=None,
            sentence_index=idx,
            severity=severity,
            fact_type=fact_type,
            suggested_source=suggested_source
        ))

    # Calculate stats
    supported = sum(1 for a in assertions if a.has_cite)
    unsupported = len(assertions) - supported
    coverage = (supported / len(assertions) * 100) if assertions else 100.0

    # Count by fact type
    material_unsupported = sum(1 for a in assertions if not a.has_cite and a.fact_type == "material")
    background_unsupported = sum(1 for a in assertions if not a.has_cite and a.fact_type == "background")

    return RecordValidation(
        total_assertions=len(assertions),
        supported_assertions=supported,
        unsupported_assertions=unsupported,
        all_cites=all_cites,
        missing_cite_flags=[a for a in assertions if not a.has_cite],
        coverage_percentage=coverage,
        cite_problems=cite_problems,
        material_facts_unsupported=material_unsupported,
        background_facts_unsupported=background_unsupported
    )


def _classify_fact_type(sentence: str) -> str:
    """Classify whether a fact is material, background, or procedural."""
    sentence_lower = sentence.lower()

    # Procedural facts
    procedural_indicators = [
        'filed', 'motion', 'court', 'ruling', 'order', 'appeal',
        'jurisdiction', 'venue', 'discovery', 'summary judgment'
    ]
    if any(ind in sentence_lower for ind in procedural_indicators):
        return "procedural"

    # Background facts (general context, don't need cites)
    background_indicators = [
        'is a corporation', 'is a company', 'is located in', 'was founded',
        'employs', 'headquartered', 'business of', 'industry'
    ]
    if any(ind in sentence_lower for ind in background_indicators):
        return "background"

    # Material facts (key disputed/determinative facts)
    material_indicators = [
        'knew', 'intended', 'caused', 'resulted', 'testified', 'admitted',
        'stated', 'wrote', 'sent', 'received', 'agreed', 'promised',
        'threatened', 'refused', 'failed', 'breached', 'violated',
        'plaintiff', 'defendant', 'on or about', 'approximately'
    ]
    if any(ind in sentence_lower for ind in material_indicators):
        return "material"

    return "material"  # Default to material for safety


def _suggest_record_source(sentence: str) -> Optional[str]:
    """Suggest which record document likely contains a fact."""
    sentence_lower = sentence.lower()

    # Testimony/statements -> Deposition or transcript
    if any(word in sentence_lower for word in ['testified', 'stated', 'said', 'told', 'admitted']):
        # Try to identify who
        name_match = re.search(r'\b([A-Z][a-z]+)\s+(?:testified|stated|said|admitted)', sentence)
        if name_match:
            return f"Check {name_match.group(1)} Deposition or Trial Transcript"
        return "Check relevant Deposition or Trial Transcript"

    # Written communications -> Exhibits
    if any(word in sentence_lower for word in ['email', 'letter', 'memo', 'wrote', 'sent']):
        return "Check Exhibits (emails/correspondence)"

    # Documents/agreements -> Exhibits
    if any(word in sentence_lower for word in ['contract', 'agreement', 'document', 'report']):
        return "Check Exhibits (contracts/documents)"

    # Financial facts -> Financial exhibits or discovery
    if any(word in sentence_lower for word in ['$', 'dollar', 'paid', 'cost', 'price', 'revenue']):
        return "Check Financial Exhibits or Interrogatory Responses"

    # Meeting/conversation facts -> Depositions
    if any(word in sentence_lower for word in ['meeting', 'conversation', 'discussed', 'call']):
        return "Check Depositions of meeting participants"

    # Employment facts -> HR documents or declarations
    if any(word in sentence_lower for word in ['hired', 'fired', 'promoted', 'terminated', 'employed']):
        return "Check HR records, Employment Agreement, or relevant Declaration"

    # Medical facts -> Medical records
    if any(word in sentence_lower for word in ['injury', 'hospital', 'doctor', 'treatment', 'diagnosis']):
        return "Check Medical Records Exhibit"

    # Dates/events -> Look for primary source
    if re.search(r'\b(?:on|in)\s+(?:January|February|March|April|May|June|July|August|September|October|November|December)', sentence):
        return "Check Exhibits or Depositions for dated events"

    return "Check Record - source unclear"


def _detect_pleading_cite_problems(text: str) -> List[CiteProblem]:
    """Detect improper citations to pleadings as proof of facts."""
    problems = []

    # Patterns for pleading citations used as proof
    pleading_patterns = [
        (r'\(Compl\.\s*(?:¶|para\.)?\s*\d+\)', 'Complaint'),
        (r'\(Am\.\s*Compl\.\s*(?:¶|para\.)?\s*\d+\)', 'Amended Complaint'),
        (r'\(Answer\s*(?:¶|para\.)?\s*\d+\)', 'Answer'),
        (r'\(Pet\.\s*(?:¶|para\.)?\s*\d+\)', 'Petition'),
    ]

    for pattern, pleading_type in pleading_patterns:
        for match in re.finditer(pattern, text, re.IGNORECASE):
            # Get context around the cite
            start = max(0, match.start() - 100)
            end = min(len(text), match.end() + 20)
            context = text[start:end]

            # Check if used for factual assertion (not procedural)
            procedural_words = ['filed', 'alleged', 'asserted', 'claims', 'pled']
            is_procedural_use = any(word in context.lower() for word in procedural_words)

            if not is_procedural_use:
                problems.append(CiteProblem(
                    cite_text=match.group(0),
                    problem_type="pleading_cite",
                    severity="error",
                    explanation=f"Citing {pleading_type} as proof of facts. Pleadings are allegations, not evidence.",
                    suggestion=f"Replace with citation to evidence (deposition, exhibit, declaration) that proves this fact"
                ))

    return problems


def _detect_vague_citations(text: str) -> List[CiteProblem]:
    """Detect vague or overly broad record citations."""
    problems = []

    # Pattern for broad page ranges (more than 10 pages)
    range_pattern = r'(?:Dep\.|Tr\.|Doc\.)\s*(?:at\s*)?(\d+)\s*-\s*(\d+)'
    for match in re.finditer(range_pattern, text, re.IGNORECASE):
        start_page = int(match.group(1))
        end_page = int(match.group(2))
        page_range = end_page - start_page

        if page_range > 25:
            problems.append(CiteProblem(
                cite_text=match.group(0),
                problem_type="vague_range",
                severity="warning",
                explanation=f"Citation spans {page_range} pages - too broad to be useful",
                suggestion="Narrow to specific page(s) where fact appears"
            ))
        elif page_range > 10:
            problems.append(CiteProblem(
                cite_text=match.group(0),
                problem_type="vague_range",
                severity="info",
                explanation=f"Citation spans {page_range} pages - consider narrowing",
                suggestion="Provide more specific pinpoint citation"
            ))

    # Pattern for deposition without line number
    dep_no_line = r'([A-Z][a-z]+)\s+Dep\.\s*(?:at\s*)?(\d+)\s*(?:[,\.\)]|$)'
    for match in re.finditer(dep_no_line, text):
        # Check if there's no line number
        if not re.search(r':\d+', match.group(0)):
            problems.append(CiteProblem(
                cite_text=match.group(0),
                problem_type="missing_pinpoint",
                severity="info",
                explanation="Deposition cite lacks line number",
                suggestion=f"Add line number: {match.group(1)} Dep. {match.group(2)}:XX"
            ))

    # Pattern for "generally" or "see generally" with record cite
    general_pattern = r'(?:see\s+)?generally\s+(?:\([^)]+\)|[A-Z][a-z]+\s+Dep\.)'
    for match in re.finditer(general_pattern, text, re.IGNORECASE):
        problems.append(CiteProblem(
            cite_text=match.group(0),
            problem_type="vague_range",
            severity="warning",
            explanation="'Generally' citation to record is too vague",
            suggestion="Provide specific pinpoint citation to support assertion"
        ))

    return problems


def format_record_validation(validation: RecordValidation) -> str:
    """Format validation results for output."""
    output = []
    output.append("RECORD CITATION VALIDATION")
    output.append("=" * 70)

    # Summary
    output.append(f"\nCoverage: {validation.coverage_percentage:.1f}%")
    output.append(f"Total factual assertions: {validation.total_assertions}")
    output.append(f"With record support: {validation.supported_assertions}")
    output.append(f"MISSING SUPPORT: {validation.unsupported_assertions}")
    output.append(f"  - Material facts unsupported: {validation.material_facts_unsupported}")
    output.append(f"  - Background facts unsupported: {validation.background_facts_unsupported}")

    # Citation problems (pleading cites, vague ranges, etc.)
    if validation.cite_problems:
        errors = [p for p in validation.cite_problems if p.severity == "error"]
        warnings = [p for p in validation.cite_problems if p.severity == "warning"]

        if errors:
            output.append(f"\n\n[!!] CITATION ERRORS ({len(errors)}):")
            output.append("-" * 40)
            for prob in errors:
                output.append(f"  {prob.cite_text}")
                output.append(f"    Problem: {prob.explanation}")
                output.append(f"    Fix: {prob.suggestion}")

        if warnings:
            output.append(f"\n\n[!] CITATION WARNINGS ({len(warnings)}):")
            output.append("-" * 40)
            for prob in warnings[:5]:
                output.append(f"  {prob.cite_text}")
                output.append(f"    {prob.explanation}")
                output.append(f"    Suggestion: {prob.suggestion}")

    # List all cites found
    if validation.all_cites:
        output.append(f"\n\nRECORD CITES FOUND ({len(validation.all_cites)}):")
        output.append("-" * 40)
        for cite in validation.all_cites[:20]:  # Limit output
            output.append(f"  {cite.raw_text} ({cite.doc_type})")

    # Flag missing cites with suggested sources
    if validation.missing_cite_flags:
        output.append(f"\n\n[!] ASSERTIONS NEEDING RECORD SUPPORT ({len(validation.missing_cite_flags)}):")
        output.append("-" * 40)

        # Sort by severity and fact type
        material_high = [a for a in validation.missing_cite_flags if a.severity == "high" and a.fact_type == "material"]
        material_other = [a for a in validation.missing_cite_flags if a.severity in ["medium", "low"] and a.fact_type == "material"]
        background = [a for a in validation.missing_cite_flags if a.fact_type == "background"]

        if material_high:
            output.append("\n[!!] CRITICAL - Material Facts (high priority):")
            for a in material_high[:5]:
                output.append(f'  "{a.text[:100]}..."')
                if a.suggested_source:
                    output.append(f"    --> {a.suggested_source}")
                else:
                    output.append("    --> [NEED CITE: ___]")

        if material_other:
            output.append(f"\n[!] Material Facts ({len(material_other)} items):")
            for a in material_other[:5]:
                output.append(f'  "{a.text[:100]}..."')
                if a.suggested_source:
                    output.append(f"    --> {a.suggested_source}")

        if background:
            output.append(f"\n[.] Background Facts (may not need cites): {len(background)} items")
            for a in background[:3]:
                output.append(f'  "{a.text[:80]}..."')

    else:
        output.append("\n[OK] All factual assertions appear to have record support!")

    return "\n".join(output)


# ============================================================================
# NEGATIVE TREATMENT DETECTION
# ============================================================================

def detect_negative_treatment(
    case_text: str,
    target_case_name: str,
    target_citation: Optional[str] = None
) -> List[NegativeTreatment]:
    """
    Detect if a case text contains negative treatment of a target case.

    Args:
        case_text: Text of the case doing the citing
        target_case_name: Name of the case to check for negative treatment
        target_citation: Optional citation to also search for

    Returns:
        List of negative treatments found
    """
    treatments = []

    # Build search terms
    search_terms = [target_case_name]
    if target_citation:
        search_terms.append(target_citation)

    # Also try short form
    if ' v. ' in target_case_name:
        short = target_case_name.split(' v. ')[0]
        search_terms.append(short)

    # Find all mentions of the target case
    for term in search_terms:
        # Escape for regex
        escaped = re.escape(term)

        for match in re.finditer(escaped, case_text, re.IGNORECASE):
            # Get context around the mention (300 chars each side)
            start = max(0, match.start() - 300)
            end = min(len(case_text), match.end() + 300)
            context = case_text[start:end]

            # Check for negative treatment patterns
            for signal, patterns in COMPILED_NEGATIVE_PATTERNS.items():
                for pattern in patterns:
                    if pattern.search(context):
                        # Extract the specific negative language
                        neg_match = pattern.search(context)
                        if neg_match:
                            # Get sentence containing the negative language
                            sent_start = context.rfind('.', 0, neg_match.start()) + 1
                            sent_end = context.find('.', neg_match.end())
                            if sent_end == -1:
                                sent_end = len(context)

                            quote = context[sent_start:sent_end].strip()

                            treatments.append(NegativeTreatment(
                                citing_case="[Current case]",
                                citing_citation="",
                                signal=signal,
                                context_quote=quote[:200],
                                severity=TREATMENT_SEVERITY.get(signal, "moderate"),
                                explanation=get_treatment_explanation(signal)
                            ))
                        break  # One match per signal per mention is enough

    # Deduplicate
    seen_quotes = set()
    unique = []
    for t in treatments:
        if t.context_quote not in seen_quotes:
            seen_quotes.add(t.context_quote)
            unique.append(t)

    return unique


def get_treatment_explanation(signal: TreatmentSignal) -> str:
    """Get explanation of what negative treatment means."""
    explanations = {
        TreatmentSignal.OVERRULED: "Case is no longer good law. DO NOT CITE.",
        TreatmentSignal.SUPERSEDED: "Statute/regulation has replaced the holding. Check current law.",
        TreatmentSignal.LIMITED: "Holding narrowed - may still apply but with restrictions.",
        TreatmentSignal.QUESTIONED: "Validity questioned - cite with caution, acknowledge criticism.",
        TreatmentSignal.CRITICIZED: "Reasoning criticized - still citable but weakened.",
        TreatmentSignal.DISTINGUISHED: "Facts distinguished - still good law for different facts.",
    }
    return explanations.get(signal, "Verify case is still good law.")


def scan_for_negative_treatment(
    cases_to_check: Dict[str, str],  # case_name -> citation
    source_texts: Dict[str, str]     # source_name -> text to scan
) -> Dict[str, List[NegativeTreatment]]:
    """
    Scan multiple source texts for negative treatment of multiple target cases.

    Args:
        cases_to_check: Cases you want to cite (name -> citation)
        source_texts: Texts to scan for negative treatment

    Returns:
        Dict mapping case name -> list of negative treatments
    """
    results = defaultdict(list)

    for case_name, citation in cases_to_check.items():
        for source_name, source_text in source_texts.items():
            treatments = detect_negative_treatment(source_text, case_name, citation)

            for t in treatments:
                t.citing_case = source_name
                results[case_name].append(t)

    return dict(results)


def format_negative_treatment(
    treatments: Dict[str, List[NegativeTreatment]]
) -> str:
    """Format negative treatment results for output."""
    output = []
    output.append("NEGATIVE TREATMENT ANALYSIS")
    output.append("=" * 70)

    if not any(treatments.values()):
        output.append("\n[OK] No negative treatment detected for checked cases.")
        output.append("(Note: This is not a substitute for full Shepardizing)")
        return "\n".join(output)

    # Group by severity
    critical = []
    serious = []
    moderate = []
    minor = []

    for case_name, case_treatments in treatments.items():
        for t in case_treatments:
            entry = (case_name, t)
            if t.severity == "critical":
                critical.append(entry)
            elif t.severity == "serious":
                serious.append(entry)
            elif t.severity == "moderate":
                moderate.append(entry)
            else:
                minor.append(entry)

    if critical:
        output.append("\n[!!] CRITICAL - DO NOT CITE WITHOUT VERIFICATION:")
        output.append("-" * 50)
        for case_name, t in critical:
            output.append(f"\n  {case_name}")
            output.append(f"  Signal: {t.signal.value.upper()}")
            output.append(f'  Context: "{t.context_quote}"')
            output.append(f"  [!]  {t.explanation}")

    if serious:
        output.append("\n[!] SERIOUS - VERIFY CURRENT SCOPE:")
        output.append("-" * 50)
        for case_name, t in serious:
            output.append(f"\n  {case_name}")
            output.append(f"  Signal: {t.signal.value}")
            output.append(f'  Context: "{t.context_quote[:100]}..."')
            output.append(f"  Note: {t.explanation}")

    if moderate:
        output.append(f"\n[!] MODERATE - CITE WITH AWARENESS: {len(moderate)} instance(s)")

    if minor:
        output.append(f"\n[.] MINOR (distinguishable): {len(minor)} instance(s)")

    output.append("\n" + "-" * 70)
    output.append("[!]  This analysis is NOT a substitute for full Shepard's/KeyCite.")
    output.append("Always verify critical authorities through official citators.")

    return "\n".join(output)


# ============================================================================
# AUTHORITY STRENGTH SCORING
# ============================================================================

def score_authority_strength(
    case_name: str,
    citation: str,
    your_jurisdiction: str,
    issue: str,
    case_text: Optional[str] = None,
    known_treatments: Optional[List[NegativeTreatment]] = None
) -> AuthorityScore:
    """
    Score the strength of a case as authority.

    Args:
        case_name: Name of the case
        citation: Full citation
        your_jurisdiction: Your circuit/court (e.g., "5th Circuit")
        issue: The legal issue
        case_text: Optional full text for additional analysis
        known_treatments: Optional known negative treatments

    Returns:
        AuthorityScore with breakdown
    """
    factors = []
    total = 0

    # 1. Court Level Score (0-40 points)
    court_score = 0

    if re.search(r'\d+\s+U\.?S\.?\s+\d+', citation) or re.search(r'\d+\s+S\.?\s*Ct\.', citation):
        court_score = 40
        factors.append("[+] Supreme Court (40 pts)")
    elif re.search(r'F\.?\s*(?:2d|3d|4th)', citation):
        court_score = 30
        factors.append("[+] Federal Circuit Court (30 pts)")
    elif re.search(r'F\.?\s*Supp', citation):
        court_score = 15
        factors.append("[ ] District Court (15 pts)")
    else:
        court_score = 10
        factors.append("[ ] State/Other Court (10 pts)")

    total += court_score

    # 2. Jurisdiction Match Score (0-25 points)
    jurisdiction_score = 0

    if court_score == 40:
        # Supreme Court is binding everywhere
        jurisdiction_score = 25
        factors.append("[+] Supreme Court - binding (25 pts)")
    else:
        # Check circuit match
        circuit_match = re.search(r'\((\d+)(?:st|nd|rd|th)\s+Cir', citation)
        dc_match = re.search(r'\(D\.?C\.?\s+Cir', citation)

        if circuit_match:
            case_circuit = f"{circuit_match.group(1)}th Circuit"
        elif dc_match:
            case_circuit = "D.C. Circuit"
        else:
            case_circuit = None

        if case_circuit and case_circuit.lower() in your_jurisdiction.lower():
            jurisdiction_score = 25
            factors.append(f"[+] Same circuit ({case_circuit}) - binding (25 pts)")
        elif case_circuit:
            jurisdiction_score = 10
            factors.append(f"[ ] Different circuit ({case_circuit}) - persuasive (10 pts)")
        else:
            jurisdiction_score = 5
            factors.append("[ ] Jurisdiction unclear (5 pts)")

    total += jurisdiction_score

    # 3. Recency Score (0-20 points)
    recency_score = 0
    year_match = re.search(r'\((\d{4})\)', citation)

    if year_match:
        year = int(year_match.group(1))
        current_year = 2024

        if year >= current_year - 3:
            recency_score = 20
            factors.append(f"[+] Recent ({year}) (20 pts)")
        elif year >= current_year - 10:
            recency_score = 15
            factors.append(f"[ ] Moderately recent ({year}) (15 pts)")
        elif year >= current_year - 20:
            recency_score = 10
            factors.append(f"[ ] Older ({year}) (10 pts)")
        else:
            recency_score = 5
            factors.append(f"[~] Old ({year}) - verify still good law (5 pts)")
    else:
        recency_score = 5
        factors.append("[ ] Year unknown (5 pts)")

    total += recency_score

    # 4. Treatment Score (0-15 points, can go negative)
    treatment_score = 15  # Start with full points

    if known_treatments:
        for t in known_treatments:
            if t.severity == "critical":
                treatment_score = -20  # Negative!
                factors.append("[-] CRITICAL negative treatment (-20 pts)")
                break
            elif t.severity == "serious":
                treatment_score -= 10
                factors.append("[-] Serious negative treatment (-10 pts)")
            elif t.severity == "moderate":
                treatment_score -= 5
                factors.append("[~] Moderate negative treatment (-5 pts)")

        if treatment_score == 15:
            factors.append("[+] No significant negative treatment (15 pts)")
    else:
        factors.append("[ ] Treatment not checked (15 pts assumed)")

    total += treatment_score

    # Ensure total is in valid range
    total = max(0, min(100, total))

    # Generate recommendation
    if total >= 80:
        recommendation = "LEAD CITATION - use first"
    elif total >= 60:
        recommendation = "Strong support - cite prominently"
    elif total >= 40:
        recommendation = "Useful support - include in string cite"
    elif total >= 20:
        recommendation = "Weak authority - use only if necessary"
    else:
        recommendation = "AVOID - find better authority"

    return AuthorityScore(
        case_name=case_name,
        citation=citation,
        total_score=total,
        court_level_score=court_score,
        jurisdiction_score=jurisdiction_score,
        recency_score=recency_score,
        treatment_score=treatment_score,
        recommendation=recommendation,
        factors=factors
    )


def rank_authorities(
    cases: Dict[str, str],  # case_name -> citation
    your_jurisdiction: str,
    issue: str
) -> List[AuthorityScore]:
    """
    Rank multiple cases by authority strength.

    Returns:
        Sorted list of AuthorityScore (highest first)
    """
    scores = []

    for case_name, citation in cases.items():
        score = score_authority_strength(case_name, citation, your_jurisdiction, issue)
        scores.append(score)

    # Sort by total score descending
    scores.sort(key=lambda x: x.total_score, reverse=True)

    return scores


def format_authority_ranking(scores: List[AuthorityScore]) -> str:
    """Format authority ranking for output."""
    output = []
    output.append("AUTHORITY STRENGTH RANKING")
    output.append("=" * 70)

    for i, score in enumerate(scores, 1):
        output.append(f"\n{i}. {score.case_name}")
        output.append(f"   {score.citation}")
        output.append(f"   SCORE: {score.total_score}/100 - {score.recommendation}")
        output.append("   Breakdown:")
        for factor in score.factors:
            output.append(f"     {factor}")

    # Citation order recommendation
    output.append("\n" + "=" * 70)
    output.append("RECOMMENDED CITATION ORDER:")

    leads = [s for s in scores if s.total_score >= 80]
    strong = [s for s in scores if 60 <= s.total_score < 80]
    support = [s for s in scores if 40 <= s.total_score < 60]

    if leads:
        output.append(f"  Lead with: {leads[0].case_name}")
    if strong:
        output.append(f"  Follow with: {', '.join(s.case_name for s in strong)}")
    if support:
        output.append(f"  String cite: {', '.join(s.case_name for s in support)}")

    return "\n".join(output)


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

def validate_record_cites_mcp(brief_text: str) -> str:
    """MCP entry point for record cite validation."""
    validation = validate_record_cites(brief_text)
    return format_record_validation(validation)


def detect_negative_treatment_mcp(
    case_name: str,
    case_citation: str,
    source_text: str
) -> str:
    """MCP entry point for negative treatment detection."""
    treatments = detect_negative_treatment(source_text, case_name, case_citation)

    if treatments:
        result = {case_name: treatments}
        return format_negative_treatment(result)
    else:
        return f"No negative treatment of {case_name} detected in the provided text.\n\n[!] This is not a substitute for full Shepardizing."


def score_authority_mcp(
    case_name: str,
    citation: str,
    your_jurisdiction: str,
    issue: str
) -> str:
    """MCP entry point for single authority scoring."""
    score = score_authority_strength(case_name, citation, your_jurisdiction, issue)

    output = []
    output.append("AUTHORITY STRENGTH ANALYSIS")
    output.append("=" * 70)
    output.append(f"\n{score.case_name}")
    output.append(f"{score.citation}")
    output.append(f"\nSCORE: {score.total_score}/100")
    output.append(f"RECOMMENDATION: {score.recommendation}")
    output.append("\nFactor Breakdown:")
    for factor in score.factors:
        output.append(f"  {factor}")

    return "\n".join(output)


def rank_authorities_mcp(
    cases: Dict[str, str],
    your_jurisdiction: str,
    issue: str
) -> str:
    """MCP entry point for ranking multiple authorities."""
    scores = rank_authorities(cases, your_jurisdiction, issue)
    return format_authority_ranking(scores)


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

if __name__ == "__main__":
    # Test record validation
    test_brief = """
    On January 15, 2023, Defendant sent an email to investors stating that
    revenue projections remained strong. (R. 145). The CEO testified at
    deposition that he knew the projections were false. Smith Dep. 45:10-20.

    The company's stock price dropped 40% after the truth was revealed.

    Plaintiff suffered losses exceeding $2 million as a result of Defendant's
    misrepresentations.
    """

    print("=" * 70)
    print("RECORD CITE VALIDATION TEST")
    print("=" * 70)
    print(validate_record_cites_mcp(test_brief))

    print("\n" * 2)
    print("=" * 70)
    print("AUTHORITY RANKING TEST")
    print("=" * 70)

    test_cases = {
        "Tellabs v. Makor": "551 U.S. 308 (2007)",
        "City of Livonia v. Boeing": "711 F.3d 754 (7th Cir. 2013)",
        "In re Enron": "761 F.3d 504 (5th Cir. 2014)",
        "Some District Case": "456 F. Supp. 2d 789 (S.D. Tex. 2018)",
    }

    print(rank_authorities_mcp(test_cases, "5th Circuit", "scienter pleading"))
