"""
Batch Processing Module
=======================
Analyze multiple legal documents simultaneously with:
- Cross-reference citation tracking across documents
- Master citation tables with frequency analysis
- Shared authority detection
- Comprehensive document comparison
"""

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

# Import our extraction functions
from text_extraction_core import (
    extract_citations_from_text,
    extract_statutes_from_text,
    normalize_citation
)

try:
    from enhanced_citations import extract_enhanced_citations, EnhancedCitation
    ENHANCED_AVAILABLE = True
except ImportError:
    ENHANCED_AVAILABLE = False


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

@dataclass
class DocumentAnalysis:
    """Analysis results for a single document."""
    doc_name: str
    citations: Dict[str, Dict]  # citation -> {pinpoint_cites, etc.}
    statutes: Dict[str, int]    # statute -> count
    word_count: int
    citation_count: int
    statute_count: int


@dataclass
class CrossReference:
    """A citation that appears in multiple documents."""
    citation: str
    documents: List[str]        # List of doc names containing this citation
    total_occurrences: int      # Total times cited across all docs
    pinpoints_by_doc: Dict[str, List[str]]  # doc_name -> pinpoint pages


@dataclass
class BatchAnalysisResult:
    """Complete batch analysis results."""
    documents: Dict[str, DocumentAnalysis]
    cross_references: List[CrossReference]
    unique_citations: Dict[str, List[str]]  # citation -> [docs containing it]
    unique_statutes: Dict[str, List[str]]   # statute -> [docs containing it]
    shared_authorities: List[Tuple[str, int, List[str]]]  # (citation, doc_count, doc_names)
    citation_frequency: Dict[str, int]      # citation -> total occurrences
    total_documents: int
    total_unique_citations: int
    total_unique_statutes: int


# ============================================================================
# BATCH ANALYSIS FUNCTIONS
# ============================================================================

def analyze_single_document(doc_name: str, text: str) -> DocumentAnalysis:
    """
    Analyze a single document and return structured results.

    Args:
        doc_name: Name/identifier for the document
        text: Document text content

    Returns:
        DocumentAnalysis with all extracted information
    """
    citations = extract_citations_from_text(text)
    statutes = extract_statutes_from_text(text)
    word_count = len(text.split())

    return DocumentAnalysis(
        doc_name=doc_name,
        citations=citations,
        statutes=statutes,
        word_count=word_count,
        citation_count=len(citations),
        statute_count=len(statutes)
    )


def analyze_batch(documents: Dict[str, str]) -> BatchAnalysisResult:
    """
    Analyze multiple documents and build cross-references.

    Args:
        documents: Dictionary mapping doc_name -> document text

    Returns:
        BatchAnalysisResult with comprehensive analysis
    """
    # Analyze each document
    doc_analyses = {}
    for doc_name, text in documents.items():
        doc_analyses[doc_name] = analyze_single_document(doc_name, text)

    # Build citation index: citation -> [docs containing it]
    citation_to_docs: Dict[str, List[str]] = defaultdict(list)
    citation_frequency: Dict[str, int] = defaultdict(int)
    citation_pinpoints: Dict[str, Dict[str, List[str]]] = defaultdict(lambda: defaultdict(list))

    for doc_name, analysis in doc_analyses.items():
        for citation, data in analysis.citations.items():
            citation_to_docs[citation].append(doc_name)
            citation_frequency[citation] += 1
            if data.get('pinpoint_cites'):
                citation_pinpoints[citation][doc_name] = data['pinpoint_cites']

    # Build statute index: statute -> [docs containing it]
    statute_to_docs: Dict[str, List[str]] = defaultdict(list)

    for doc_name, analysis in doc_analyses.items():
        for statute in analysis.statutes.keys():
            statute_to_docs[statute].append(doc_name)

    # Build cross-references (citations appearing in multiple docs)
    cross_refs = []
    for citation, doc_list in citation_to_docs.items():
        if len(doc_list) > 1:  # Appears in multiple documents
            cross_refs.append(CrossReference(
                citation=citation,
                documents=doc_list,
                total_occurrences=citation_frequency[citation],
                pinpoints_by_doc=dict(citation_pinpoints[citation])
            ))

    # Sort cross-references by number of documents (most shared first)
    cross_refs.sort(key=lambda x: len(x.documents), reverse=True)

    # Build shared authorities list
    shared_authorities = [
        (citation, len(docs), docs)
        for citation, docs in citation_to_docs.items()
        if len(docs) > 1
    ]
    shared_authorities.sort(key=lambda x: x[1], reverse=True)

    return BatchAnalysisResult(
        documents=doc_analyses,
        cross_references=cross_refs,
        unique_citations=dict(citation_to_docs),
        unique_statutes=dict(statute_to_docs),
        shared_authorities=shared_authorities,
        citation_frequency=dict(citation_frequency),
        total_documents=len(documents),
        total_unique_citations=len(citation_to_docs),
        total_unique_statutes=len(statute_to_docs)
    )


# ============================================================================
# CITATION TABLE GENERATION
# ============================================================================

def generate_master_citation_table(result: BatchAnalysisResult) -> str:
    """
    Generate a comprehensive citation table from batch analysis.

    Args:
        result: BatchAnalysisResult from analyze_batch()

    Returns:
        Formatted citation table string
    """
    output = []
    output.append("=" * 80)
    output.append("MASTER CITATION TABLE")
    output.append("=" * 80)
    output.append("")

    # Summary statistics
    output.append("SUMMARY")
    output.append("-" * 80)
    output.append(f"Documents analyzed: {result.total_documents}")
    output.append(f"Unique case citations: {result.total_unique_citations}")
    output.append(f"Unique statutes/rules: {result.total_unique_statutes}")
    output.append(f"Shared authorities (in 2+ docs): {len(result.shared_authorities)}")
    output.append("")

    # Document breakdown
    output.append("DOCUMENT BREAKDOWN")
    output.append("-" * 80)
    for doc_name, analysis in result.documents.items():
        output.append(f"\n{doc_name}:")
        output.append(f"  Words: {analysis.word_count:,}")
        output.append(f"  Case citations: {analysis.citation_count}")
        output.append(f"  Statutes/rules: {analysis.statute_count}")
    output.append("")

    # Shared authorities (most important)
    if result.shared_authorities:
        output.append("")
        output.append("SHARED AUTHORITIES (Citations in Multiple Documents)")
        output.append("-" * 80)
        output.append("")

        for citation, doc_count, doc_names in result.shared_authorities[:20]:  # Top 20
            output.append(f"• {citation}")
            output.append(f"  Appears in {doc_count} documents: {', '.join(doc_names)}")

            # Show pinpoints if available
            for cross_ref in result.cross_references:
                if cross_ref.citation == citation and cross_ref.pinpoints_by_doc:
                    for doc, pinpoints in cross_ref.pinpoints_by_doc.items():
                        if pinpoints:
                            output.append(f"    {doc} cites pages: {', '.join(pinpoints)}")
            output.append("")

    # All citations by frequency
    output.append("")
    output.append("ALL CITATIONS BY FREQUENCY")
    output.append("-" * 80)
    output.append("")

    sorted_citations = sorted(
        result.citation_frequency.items(),
        key=lambda x: (len(result.unique_citations.get(x[0], [])), x[1]),
        reverse=True
    )

    for citation, freq in sorted_citations:
        docs = result.unique_citations.get(citation, [])
        if len(docs) > 1:
            marker = f"[{len(docs)} docs]"
        else:
            marker = f"[{docs[0]}]"
        output.append(f"  {citation} {marker}")

    # Statutes section
    if result.unique_statutes:
        output.append("")
        output.append("")
        output.append("STATUTES AND RULES")
        output.append("-" * 80)
        output.append("")

        # Shared statutes first
        shared_statutes = [(s, docs) for s, docs in result.unique_statutes.items() if len(docs) > 1]
        if shared_statutes:
            output.append("Shared across documents:")
            for statute, docs in sorted(shared_statutes, key=lambda x: len(x[1]), reverse=True):
                output.append(f"  • {statute}")
                output.append(f"    In: {', '.join(docs)}")
            output.append("")

        # All statutes
        output.append("All statutes/rules:")
        for statute, docs in sorted(result.unique_statutes.items()):
            if len(docs) > 1:
                output.append(f"  • {statute} [{len(docs)} docs]")
            else:
                output.append(f"  • {statute} [{docs[0]}]")

    return '\n'.join(output)


def generate_document_comparison(result: BatchAnalysisResult) -> str:
    """
    Generate a comparison showing overlap between documents.

    Args:
        result: BatchAnalysisResult from analyze_batch()

    Returns:
        Formatted comparison string
    """
    output = []
    output.append("=" * 80)
    output.append("DOCUMENT COMPARISON - CITATION OVERLAP")
    output.append("=" * 80)
    output.append("")

    doc_names = list(result.documents.keys())

    if len(doc_names) < 2:
        output.append("Need at least 2 documents for comparison.")
        return '\n'.join(output)

    # Pairwise comparison
    for i, doc_a in enumerate(doc_names):
        for doc_b in doc_names[i+1:]:
            cites_a = set(result.documents[doc_a].citations.keys())
            cites_b = set(result.documents[doc_b].citations.keys())

            shared = cites_a & cites_b
            only_a = cites_a - cites_b
            only_b = cites_b - cites_a

            output.append(f"\n{doc_a} vs {doc_b}")
            output.append("-" * 60)
            output.append(f"Shared citations: {len(shared)}")
            output.append(f"Only in {doc_a}: {len(only_a)}")
            output.append(f"Only in {doc_b}: {len(only_b)}")

            if shared:
                # Calculate overlap percentage
                overlap_pct = (len(shared) / min(len(cites_a), len(cites_b))) * 100 if min(len(cites_a), len(cites_b)) > 0 else 0
                output.append(f"Overlap: {overlap_pct:.1f}%")
                output.append("")
                output.append("Shared citations:")
                for cite in sorted(shared):
                    output.append(f"  • {cite}")

    return '\n'.join(output)


def generate_authority_matrix(result: BatchAnalysisResult) -> str:
    """
    Generate a matrix showing which authorities appear in which documents.

    Args:
        result: BatchAnalysisResult from analyze_batch()

    Returns:
        Formatted matrix string
    """
    output = []
    output.append("=" * 80)
    output.append("AUTHORITY MATRIX")
    output.append("=" * 80)
    output.append("")

    doc_names = list(result.documents.keys())

    # Header row
    output.append("Citation".ljust(50) + " | " + " | ".join(d[:10].center(10) for d in doc_names))
    output.append("-" * (52 + 13 * len(doc_names)))

    # Sort citations by how many docs contain them
    sorted_citations = sorted(
        result.unique_citations.items(),
        key=lambda x: len(x[1]),
        reverse=True
    )

    for citation, docs in sorted_citations[:30]:  # Top 30
        # Truncate long citations
        cite_display = citation[:48] + ".." if len(citation) > 50 else citation.ljust(50)

        # Mark which docs contain this citation
        markers = []
        for doc_name in doc_names:
            if doc_name in docs:
                markers.append("    ✓     ")
            else:
                markers.append("          ")

        output.append(f"{cite_display} | {' | '.join(markers)}")

    if len(result.unique_citations) > 30:
        output.append(f"\n... and {len(result.unique_citations) - 30} more citations")

    return '\n'.join(output)


# ============================================================================
# FORMATTED OUTPUT FOR MCP
# ============================================================================

def batch_analyze_formatted(documents: Dict[str, str]) -> str:
    """
    Main entry point - analyze documents and return formatted report.

    Args:
        documents: Dictionary mapping doc_name -> document text

    Returns:
        Complete formatted analysis report
    """
    if not documents:
        return "No documents provided for analysis."

    if len(documents) == 1:
        return "Batch analysis requires at least 2 documents. Use single-document tools for one document."

    result = analyze_batch(documents)

    output = []
    output.append(generate_master_citation_table(result))
    output.append("")
    output.append(generate_document_comparison(result))

    # Only include matrix if not too many documents
    if len(documents) <= 5:
        output.append("")
        output.append(generate_authority_matrix(result))

    return '\n'.join(output)


def get_batch_data(documents: Dict[str, str]) -> Dict:
    """
    Get batch analysis as dictionary for programmatic use.

    Args:
        documents: Dictionary mapping doc_name -> document text

    Returns:
        Dictionary with all analysis data
    """
    result = analyze_batch(documents)

    return {
        'total_documents': result.total_documents,
        'total_unique_citations': result.total_unique_citations,
        'total_unique_statutes': result.total_unique_statutes,
        'shared_authority_count': len(result.shared_authorities),
        'documents': {
            name: {
                'word_count': analysis.word_count,
                'citation_count': analysis.citation_count,
                'statute_count': analysis.statute_count,
                'citations': list(analysis.citations.keys()),
                'statutes': list(analysis.statutes.keys())
            }
            for name, analysis in result.documents.items()
        },
        'shared_authorities': [
            {'citation': cite, 'document_count': count, 'documents': docs}
            for cite, count, docs in result.shared_authorities
        ],
        'cross_references': [
            {
                'citation': cr.citation,
                'documents': cr.documents,
                'total_occurrences': cr.total_occurrences,
                'pinpoints_by_doc': cr.pinpoints_by_doc
            }
            for cr in result.cross_references
        ],
        'all_citations': result.unique_citations,
        'all_statutes': result.unique_statutes,
        'citation_frequency': result.citation_frequency
    }


# ============================================================================
# CONVENIENCE FUNCTIONS
# ============================================================================

def find_common_authorities(documents: Dict[str, str], min_docs: int = 2) -> List[Tuple[str, int, List[str]]]:
    """
    Find citations that appear in at least min_docs documents.

    Args:
        documents: Dictionary mapping doc_name -> document text
        min_docs: Minimum number of documents a citation must appear in

    Returns:
        List of (citation, doc_count, doc_names) tuples
    """
    result = analyze_batch(documents)
    return [
        (cite, count, docs)
        for cite, count, docs in result.shared_authorities
        if count >= min_docs
    ]


def get_citation_coverage(documents: Dict[str, str], citation: str) -> Dict[str, bool]:
    """
    Check which documents contain a specific citation.

    Args:
        documents: Dictionary mapping doc_name -> document text
        citation: The citation to search for

    Returns:
        Dictionary mapping doc_name -> True/False
    """
    result = analyze_batch(documents)
    docs_with_cite = result.unique_citations.get(citation, [])

    return {
        doc_name: doc_name in docs_with_cite
        for doc_name in documents.keys()
    }
