"""
Case Bank - Build searchable index of rules/holdings from case files.

This module provides functionality to:
1. Process all case files in a directory
2. Extract rules and holdings from each case
3. Build a searchable JSON index
4. Organize rules by type with cross-references
"""

import json
import os
import re
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Any
from datetime import datetime

# Import from existing modules
try:
    from brief_synthesis import read_case_file, _extract_citation_from_text, generate_parenthetical
    from enhanced_rules import extract_enhanced_rules
except ImportError:
    # Fallback for standalone testing
    pass


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

@dataclass
class CaseBankEntry:
    """A single case in the bank."""
    case_name: str
    citation: Optional[str]
    file_path: str
    file_type: str
    rules: List[Dict]
    holdings: List[str]
    legal_areas: List[str]
    key_terms: List[str]
    parenthetical: Optional[str] = None


@dataclass
class CaseBank:
    """Searchable case bank index."""
    cases: Dict[str, Dict] = field(default_factory=dict)
    rules_by_type: Dict[str, List[Dict]] = field(default_factory=dict)
    cross_references: Dict[str, List[str]] = field(default_factory=dict)
    total_rules: int = 0
    total_cases: int = 0
    created_at: str = ""
    index_path: Optional[str] = None
    errors: List[Dict] = field(default_factory=list)


# ============================================================================
# CORE FUNCTIONS
# ============================================================================

def build_case_bank(
    directory_path: str,
    file_pattern: str = "*.pdf"
) -> str:
    """
    Process all matching files in directory, build searchable index.

    Args:
        directory_path: Directory containing case files
        file_pattern: Glob pattern (e.g., "*.pdf", "*.txt", "*.*")

    Returns:
        Formatted output with summary and JSON data
    """
    bank = CaseBank(
        created_at=datetime.now().isoformat(),
        rules_by_type={
            "holding": [],
            "test": [],
            "elements": [],
            "factors": [],
            "standard": [],
            "bright_line": [],
            "balancing": [],
            "multi_factor": [],
            "totality": [],
            "presumption": [],
            "definition": [],
            "other": []
        }
    )

    # Find matching files
    dir_path = Path(directory_path)
    if not dir_path.exists():
        return _format_bank_error(f"Directory not found: {directory_path}")

    if not dir_path.is_dir():
        return _format_bank_error(f"Path is not a directory: {directory_path}")

    matching_files = list(dir_path.glob(file_pattern))
    if not matching_files:
        return _format_bank_error(f"No files matching '{file_pattern}' in {directory_path}")

    # Process each file
    for file_path in matching_files:
        try:
            entry = _process_case_file(str(file_path))
            if entry:
                bank.cases[entry.case_name] = asdict(entry)
                bank.total_cases += 1

                # Organize rules by type
                for rule in entry.rules:
                    rule_type = rule.get("rule_type", "other").lower()
                    if rule_type not in bank.rules_by_type:
                        rule_type = "other"

                    bank.rules_by_type[rule_type].append({
                        "text": rule.get("text"),
                        "case_name": entry.case_name,
                        "citation": entry.citation,
                        "pinpoint": rule.get("pinpoint"),
                        "elements": rule.get("elements", []),
                        "factors": rule.get("factors", [])
                    })
                    bank.total_rules += 1

                # Track cross-references (cases citing other cases)
                _extract_cross_references(entry, bank)

        except Exception as e:
            bank.errors.append({
                "file": str(file_path),
                "error": str(e)
            })

    # Save index to file
    index_path = dir_path / "case_bank_index.json"
    try:
        with open(index_path, 'w', encoding='utf-8') as f:
            json.dump(_bank_to_dict(bank), f, indent=2, default=str)
        bank.index_path = str(index_path)
    except Exception as e:
        bank.errors.append({"file": str(index_path), "error": f"Failed to save: {e}"})

    return _format_bank_output(bank)


def _process_case_file(file_path: str) -> Optional[CaseBankEntry]:
    """Process a single case file into a CaseBankEntry."""
    # Read file
    text, file_type = read_case_file(file_path)
    if not text or len(text.strip()) < 100:
        return None

    # Extract case name from filename
    basename = os.path.basename(file_path)
    name_without_ext = os.path.splitext(basename)[0]
    case_name = re.sub(r'\s+v\s+', ' v. ', name_without_ext)
    case_name = re.sub(r'[_-]', ' ', case_name)

    # Extract citation (pass case_name to help find the right citation)
    citation = _extract_citation_from_text(text, case_name)

    # Extract rules
    rules = []
    holdings = []
    legal_areas = set()
    key_terms = set()

    try:
        extracted_rules = extract_enhanced_rules(text, case_name, citation)
        for r in extracted_rules:
            rule_dict = {
                "text": r.text,
                "rule_type": r.rule_type.value if hasattr(r.rule_type, 'value') else str(r.rule_type),
                "level": r.level.value if hasattr(r.level, 'value') else str(r.level),
                "pinpoint": r.pinpoint_page,
                "elements": r.elements,
                "factors": r.factors
            }
            rules.append(rule_dict)

            # Track holdings specifically
            if r.rule_type.value in ("holding", "test") if hasattr(r.rule_type, 'value') else False:
                holdings.append(r.text)

            # Collect legal areas and key terms
            if r.legal_area:
                legal_areas.add(r.legal_area)
            if r.key_terms:
                key_terms.update(r.key_terms)
    except Exception:
        pass

    # Generate parenthetical
    parenthetical = None
    try:
        paren = generate_parenthetical(case_name, text, citation or "")
        if paren:
            parenthetical = paren.parenthetical_text
    except Exception:
        pass

    return CaseBankEntry(
        case_name=case_name,
        citation=citation,
        file_path=file_path,
        file_type=file_type,
        rules=rules,
        holdings=holdings,
        legal_areas=list(legal_areas),
        key_terms=list(key_terms)[:20],  # Limit key terms
        parenthetical=parenthetical
    )


def _extract_cross_references(entry: CaseBankEntry, bank: CaseBank):
    """Extract cases referenced in this case's text."""
    # Look for case names in the rules text
    for rule in entry.rules:
        text = rule.get("text", "")
        # Find patterns like "in Smith v. Jones" or "Smith held"
        case_refs = re.findall(r'\b([A-Z][a-z]+)\s+v\.?\s+([A-Z][a-z]+)', text)
        for ref in case_refs:
            ref_name = f"{ref[0]} v. {ref[1]}"
            if ref_name != entry.case_name:
                if entry.case_name not in bank.cross_references:
                    bank.cross_references[entry.case_name] = []
                if ref_name not in bank.cross_references[entry.case_name]:
                    bank.cross_references[entry.case_name].append(ref_name)


def _bank_to_dict(bank: CaseBank) -> Dict:
    """Convert CaseBank to JSON-serializable dict."""
    return {
        "metadata": {
            "created_at": bank.created_at,
            "total_cases": bank.total_cases,
            "total_rules": bank.total_rules,
            "index_path": bank.index_path
        },
        "cases": bank.cases,
        "rules_by_type": bank.rules_by_type,
        "cross_references": bank.cross_references,
        "errors": bank.errors
    }


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

def _format_bank_output(bank: CaseBank) -> str:
    """Format case bank as text + JSON output."""
    output = []
    output.append("=" * 70)
    output.append("CASE BANK BUILD COMPLETE")
    output.append("=" * 70)

    output.append(f"\nProcessed: {bank.total_cases} cases")
    output.append(f"Total rules: {bank.total_rules}")
    if bank.index_path:
        output.append(f"Index saved: {bank.index_path}")

    # Rules by type summary
    output.append("\n--- RULES BY TYPE ---")
    for rule_type, rules in bank.rules_by_type.items():
        if rules:
            output.append(f"  {rule_type.upper()}: {len(rules)}")

    # Case list
    output.append("\n--- CASES INDEXED ---")
    for case_name, entry in list(bank.cases.items())[:20]:
        citation = entry.get("citation", "No citation")
        rule_count = len(entry.get("rules", []))
        output.append(f"  - {case_name} ({citation}) - {rule_count} rules")

    if len(bank.cases) > 20:
        output.append(f"  ... and {len(bank.cases) - 20} more")

    # Cross-references
    if bank.cross_references:
        output.append("\n--- CROSS-REFERENCES ---")
        for case_name, refs in list(bank.cross_references.items())[:10]:
            if refs:
                output.append(f"  {case_name} cites: {', '.join(refs[:3])}")

    # Errors
    if bank.errors:
        output.append("\n--- ERRORS ---")
        for err in bank.errors[:5]:
            output.append(f"  {err['file']}: {err['error']}")

    # JSON
    output.append("\n" + "=" * 70)
    output.append("JSON_DATA:")
    output.append(json.dumps(_bank_to_dict(bank), indent=2, default=str))

    return "\n".join(output)


def _format_bank_error(message: str) -> str:
    """Format an error message."""
    return f"""{"=" * 70}
CASE BANK BUILD FAILED
{"=" * 70}

ERROR: {message}

JSON_DATA:
{{"error": "{message}"}}
"""


# ============================================================================
# SEARCH FUNCTIONS (for future use)
# ============================================================================

def search_case_bank(
    index_path: str,
    query: Optional[str] = None,
    rule_type: Optional[str] = None,
    case_name: Optional[str] = None
) -> List[Dict]:
    """
    Search a case bank index.

    Args:
        index_path: Path to case_bank_index.json
        query: Text to search for in rules
        rule_type: Filter by rule type (holding, elements, etc.)
        case_name: Filter by case name

    Returns:
        List of matching rules/cases
    """
    with open(index_path, 'r', encoding='utf-8') as f:
        bank = json.load(f)

    results = []

    # Search in rules_by_type
    rules_to_search = bank.get("rules_by_type", {})

    if rule_type:
        rules_to_search = {rule_type: rules_to_search.get(rule_type, [])}

    for rtype, rules in rules_to_search.items():
        for rule in rules:
            # Apply filters
            if case_name and rule.get("case_name", "").lower() != case_name.lower():
                continue

            if query:
                text = rule.get("text", "").lower()
                if query.lower() not in text:
                    continue

            results.append({
                "rule_type": rtype,
                **rule
            })

    return results


# ============================================================================
# SEMANTIC SEARCH & QUERY EXPANSION
# ============================================================================

# Legal term expansion dictionary - maps terms to related concepts
LEGAL_TERM_EXPANSIONS = {
    # Procedural terms
    "scienter": ["fraudulent intent", "intent to deceive", "state of mind", "knowledge", "recklessness"],
    "plead": ["allege", "alleged", "plausibly alleged", "adequately alleged", "pleading"],
    "pled": ["alleged", "pleaded", "asserted", "claimed"],
    "dismiss": ["dismissed", "dismissal", "12(b)(6)", "failure to state a claim"],
    "summary judgment": ["no genuine dispute", "material fact", "reasonable jury", "Rule 56"],

    # Evidence terms
    "causation": ["cause", "caused", "causing", "proximate cause", "but-for", "substantial factor"],
    "temporal proximity": ["temporal relationship", "close in time", "timing", "shortly after", "immediately following"],
    "preponderance": ["more likely than not", "greater weight", "50%", "balance of probabilities"],
    "clear and convincing": ["highly probable", "firm belief", "reasonable certainty"],
    "beyond reasonable doubt": ["moral certainty", "firmly convinced", "no reasonable doubt"],

    # Substantive terms
    "negligence": ["duty", "breach", "causation", "damages", "reasonable care", "foreseeable"],
    "fraud": ["misrepresentation", "false statement", "material", "reliance", "damages", "scienter"],
    "breach of contract": ["material breach", "performance", "non-performance", "damages", "consideration"],
    "standing": ["injury in fact", "causation", "redressability", "concrete", "particularized"],

    # Remedies
    "damages": ["compensatory", "punitive", "nominal", "actual", "consequential", "incidental"],
    "injunction": ["equitable relief", "irreparable harm", "balance of hardships", "public interest"],

    # Standards
    "de novo": ["fresh review", "no deference", "independently"],
    "abuse of discretion": ["clearly erroneous", "arbitrary", "capricious", "unreasonable"],
    "substantial evidence": ["reasonable mind", "adequate support", "record evidence"],
}

# Sentence role keywords for semantic filtering
ROLE_KEYWORDS = {
    "holding": ["hold", "held", "holding", "conclude", "concluded", "find", "found", "determine"],
    "rule": ["must", "shall", "required", "standard", "test", "elements", "factors"],
    "finding": ["evidence shows", "record establishes", "testimony", "jury found", "factfinder"],
    "reasoning": ["because", "since", "therefore", "thus", "accordingly", "policy", "purpose"],
}


def expand_legal_query(query: str) -> Dict[str, Any]:
    """
    Expand a natural language legal query into structured search terms.

    Parses the query to extract:
    - Core concepts and their expansions
    - Desired sentence role (holding, rule, finding)
    - Relationship type (sufficiency, causation, etc.)

    Args:
        query: Natural language query like "Find holdings that scienter can be pled from reckless conduct"

    Returns:
        Structured query with expanded terms and filters
    """
    query_lower = query.lower()

    # Detect desired sentence role
    target_role = None
    for role, keywords in ROLE_KEYWORDS.items():
        if any(kw in query_lower for kw in keywords):
            target_role = role
            break

    # Also check explicit role requests
    if "holding" in query_lower:
        target_role = "holding"
    elif "rule" in query_lower or "standard" in query_lower:
        target_role = "rule"
    elif "finding" in query_lower or "evidence" in query_lower:
        target_role = "finding"

    # Extract and expand core terms
    expanded_terms = []
    matched_terms = []

    for term, expansions in LEGAL_TERM_EXPANSIONS.items():
        if term in query_lower:
            matched_terms.append(term)
            expanded_terms.extend(expansions)
            expanded_terms.append(term)

    # Remove duplicates while preserving order
    seen = set()
    unique_expanded = []
    for t in expanded_terms:
        if t.lower() not in seen:
            seen.add(t.lower())
            unique_expanded.append(t)

    # Extract quoted phrases as exact matches
    exact_phrases = re.findall(r'"([^"]+)"', query)

    # Detect relationship/analysis type
    relationship = None
    if any(word in query_lower for word in ["sufficiency", "sufficient", "adequate", "enough"]):
        relationship = "sufficiency"
    elif any(word in query_lower for word in ["causation", "cause", "caused"]):
        relationship = "causation"
    elif any(word in query_lower for word in ["element", "elements", "factor", "factors"]):
        relationship = "elements"
    elif any(word in query_lower for word in ["burden", "proof", "persuasion"]):
        relationship = "burden"

    return {
        "original_query": query,
        "target_role": target_role,
        "matched_terms": matched_terms,
        "expanded_terms": unique_expanded,
        "exact_phrases": exact_phrases,
        "relationship_type": relationship,
        "search_strategy": "semantic" if matched_terms else "keyword"
    }


def semantic_case_search(
    index_path: str,
    query: str,
    max_results: int = 20,
    min_relevance: float = 0.3
) -> Dict[str, Any]:
    """
    Semantic search of case bank using query expansion and role filtering.

    Unlike basic search, this:
    1. Expands query terms to related legal concepts
    2. Filters by sentence role (if requested)
    3. Ranks results by relevance score
    4. Returns matching sentences with context

    Args:
        index_path: Path to case_bank_index.json
        query: Natural language query
        max_results: Maximum results to return
        min_relevance: Minimum relevance score (0-1)

    Returns:
        Ranked results with relevance scores and matched terms
    """
    # Load case bank
    with open(index_path, 'r', encoding='utf-8') as f:
        bank = json.load(f)

    # Expand the query
    expanded = expand_legal_query(query)

    # Build search terms list
    search_terms = set()
    for term in expanded.get("expanded_terms", []):
        search_terms.add(term.lower())
    for phrase in expanded.get("exact_phrases", []):
        search_terms.add(phrase.lower())

    # Also add words from original query
    query_words = [w.lower() for w in query.split() if len(w) > 3]
    search_terms.update(query_words)

    # Filter by rule type if target role specified
    target_role = expanded.get("target_role")
    rule_type_filter = None
    if target_role == "holding":
        rule_type_filter = "holding"
    elif target_role == "rule":
        rule_type_filter = None  # Search all rule types

    # Score and collect results
    scored_results = []

    rules_to_search = bank.get("rules_by_type", {})
    if rule_type_filter:
        rules_to_search = {rule_type_filter: rules_to_search.get(rule_type_filter, [])}

    for rtype, rules in rules_to_search.items():
        for rule in rules:
            rule_text = rule.get("text", "").lower()

            # Calculate relevance score
            matched_count = 0
            matched_terms_in_rule = []

            for term in search_terms:
                if term in rule_text:
                    matched_count += 1
                    matched_terms_in_rule.append(term)

            if matched_count == 0:
                continue

            # Score based on match density and term importance
            relevance = matched_count / max(len(search_terms), 1)

            # Boost for exact phrase matches
            for phrase in expanded.get("exact_phrases", []):
                if phrase.lower() in rule_text:
                    relevance += 0.2

            # Boost for matched expansion terms (more specific)
            expansion_matches = sum(1 for t in expanded.get("matched_terms", [])
                                   if t.lower() in rule_text)
            relevance += expansion_matches * 0.15

            # Cap at 1.0
            relevance = min(relevance, 1.0)

            if relevance >= min_relevance:
                scored_results.append({
                    "case_name": rule.get("case_name", "Unknown"),
                    "citation": rule.get("citation"),
                    "rule_type": rtype,
                    "text": rule.get("text", ""),
                    "pinpoint": rule.get("pinpoint"),
                    "relevance_score": round(relevance, 3),
                    "matched_terms": matched_terms_in_rule[:5],
                    "elements": rule.get("elements", []),
                    "factors": rule.get("factors", [])
                })

    # Sort by relevance
    scored_results.sort(key=lambda x: x["relevance_score"], reverse=True)

    # Limit results
    top_results = scored_results[:max_results]

    return {
        "_meta": {
            "query_expansion": expanded,
            "total_matches": len(scored_results),
            "returned": len(top_results),
            "search_terms_used": list(search_terms)[:20]
        },
        "results": top_results
    }


def format_semantic_search(result: Dict[str, Any]) -> str:
    """Format semantic search results for display."""
    output = []
    output.append("=" * 70)
    output.append("SEMANTIC CASE SEARCH RESULTS")
    output.append("=" * 70)

    meta = result.get("_meta", {})
    expansion = meta.get("query_expansion", {})

    output.append(f"\nOriginal query: {expansion.get('original_query', '')}")
    output.append(f"Target role: {expansion.get('target_role', 'any')}")
    output.append(f"Matched terms: {', '.join(expansion.get('matched_terms', []))}")
    output.append(f"Search strategy: {expansion.get('search_strategy', 'keyword')}")
    output.append(f"\nTotal matches: {meta.get('total_matches', 0)}")
    output.append(f"Showing: {meta.get('returned', 0)} results")

    results = result.get("results", [])
    for i, r in enumerate(results, 1):
        output.append(f"\n--- Result {i} (Relevance: {r['relevance_score']:.0%}) ---")
        output.append(f"Case: {r['case_name']}")
        if r.get('citation'):
            output.append(f"Citation: {r['citation']}")
        output.append(f"Type: {r['rule_type'].upper()}")
        output.append(f"Matched: {', '.join(r.get('matched_terms', []))}")
        output.append(f"Rule: {r['text'][:300]}{'...' if len(r['text']) > 300 else ''}")
        if r.get('elements'):
            output.append(f"Elements: {r['elements'][:3]}")

    return "\n".join(output)


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

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        result = build_case_bank(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else "*.pdf")
        print(result)
    else:
        print("Usage: python case_bank.py <directory> [file_pattern]")
