Designing a Legal Document Search System.

# search/legal_query_parser.py

"""
Legal query parser supporting boolean and proximity search.

Applies: Week 7, Day 3 - Query Processing
"""

from dataclasses import dataclass
from typing import List, Optional, Union
from enum import Enum, auto
import re


class TokenType(Enum):
    """Token types in legal queries."""
    TERM = auto()
    PHRASE = auto()
    AND = auto()
    OR = auto()
    NOT = auto()
    LPAREN = auto()
    RPAREN = auto()
    PROXIMITY = auto()      # /5, /10, /s, /p
    FIELD = auto()          # judge:, court:, date:
    CITATION = auto()       # cites:, citedby:
    RANGE = auto()          # [2020 TO 2024]


@dataclass
class Token:
    """A lexical token."""
    type: TokenType
    value: str
    position: int


class LegalQueryLexer:
    """
    Tokenizes legal search queries.
    """
    
    KEYWORDS = {
        'AND': TokenType.AND,
        'OR': TokenType.OR,
        'NOT': TokenType.NOT,
        'AND NOT': TokenType.NOT,
    }
    
    FIELD_PREFIXES = ['judge:', 'court:', 'date:', 'party:', 'title:']
    CITATION_PREFIXES = ['cites:', 'citedby:', 'citing:']
    
    PROXIMITY_PATTERN = re.compile(r'/(\d+|s|p)')
    
    def tokenize(self, query: str) -> List[Token]:
        """Tokenize a legal query string."""
        
        tokens = []
        position = 0
        
        while position < len(query):
            # Skip whitespace
            if query[position].isspace():
                position += 1
                continue
            
            # Parentheses
            if query[position] == '(':
                tokens.append(Token(TokenType.LPAREN, '(', position))
                position += 1
                continue
            
            if query[position] == ')':
                tokens.append(Token(TokenType.RPAREN, ')', position))
                position += 1
                continue
            
            # Phrase (quoted string)
            if query[position] == '"':
                end = query.find('"', position + 1)
                if end == -1:
                    end = len(query)
                phrase = query[position+1:end]
                tokens.append(Token(TokenType.PHRASE, phrase, position))
                position = end + 1
                continue
            
            # Proximity operator
            if query[position] == '/':
                match = self.PROXIMITY_PATTERN.match(query[position:])
                if match:
                    tokens.append(Token(TokenType.PROXIMITY, match.group(1), position))
                    position += len(match.group(0))
                    continue
            
            # Range [2020 TO 2024]
            if query[position] == '[':
                end = query.find(']', position)
                if end != -1:
                    range_expr = query[position:end+1]
                    tokens.append(Token(TokenType.RANGE, range_expr, position))
                    position = end + 1
                    continue
            
            # Word or keyword
            word_match = re.match(r'[\w\.:]+', query[position:])
            if word_match:
                word = word_match.group(0)
                
                # Check for field prefix
                for prefix in self.FIELD_PREFIXES:
                    if word.lower().startswith(prefix):
                        tokens.append(Token(TokenType.FIELD, word, position))
                        break
                else:
                    # Check for citation prefix
                    for prefix in self.CITATION_PREFIXES:
                        if word.lower().startswith(prefix):
                            # Reporter citations contain spaces: "410 U.S. 113"
                            cite = re.match(r'\d+\s+[A-Za-z][\w\.]*\s+\d+',
                                            query[position + len(prefix):])
                            if cite:
                                word = prefix + cite.group(0)
                            tokens.append(Token(TokenType.CITATION, word, position))
                            break
                    else:
                        # Check for keyword
                        upper = word.upper()
                        if upper in self.KEYWORDS:
                            tokens.append(Token(self.KEYWORDS[upper], word, position))
                        else:
                            tokens.append(Token(TokenType.TERM, word, position))
                
                position += len(word)
                continue
            
            # Unknown character, skip
            position += 1
        
        return tokens


class LegalQueryBuilder:
    """
    Builds Elasticsearch queries from parsed legal queries.
    """
    
    def __init__(self):
        self.lexer = LegalQueryLexer()
    
    def build_query(
        self,
        query_string: str,
        tenant_id: Optional[str] = None,
        search_fields: List[str] = None
    ) -> dict:
        """
        Build Elasticsearch query from legal query string.
        
        Applies concepts:
        - Query vs Filter context (Day 3)
        - Multi-match for field boosting (Day 3)
        - Span queries for proximity (new in this capstone)
        """
        
        if not search_fields:
            search_fields = ["content", "title^2", "summary^1.5"]
        
        tokens = self.lexer.tokenize(query_string)
        
        # Parse tokens into ES query
        es_query = self._parse_expression(tokens, search_fields)
        
        # Wrap with filters
        final_query = {
            "query": {
                "bool": {
                    "must": [es_query],
                    "filter": []
                }
            }
        }
        
        # Add tenant isolation filter (critical for security!)
        if tenant_id:
            final_query["query"]["bool"]["filter"].append({
                "term": {"tenant_id": tenant_id}
            })
        
        return final_query
    
    def _parse_expression(
        self,
        tokens: List[Token],
        fields: List[str]
    ) -> dict:
        """Parse tokens into ES query structure."""
        
        if not tokens:
            return {"match_all": {}}
        
        # Simple case: single term
        if len(tokens) == 1:
            return self._token_to_query(tokens[0], fields)
        
        # Handle boolean operators
        must_clauses = []
        should_clauses = []
        must_not_clauses = []
        
        i = 0
        current_operator = TokenType.AND  # Default is AND
        
        while i < len(tokens):
            token = tokens[i]
            
            if token.type == TokenType.AND:
                current_operator = TokenType.AND
                i += 1
                continue
            
            if token.type == TokenType.OR:
                current_operator = TokenType.OR
                # The left operand becomes one alternative of the OR group
                # (simplified precedence: OR groups are ANDed with the rest)
                if must_clauses and not should_clauses:
                    should_clauses.append(must_clauses.pop())
                i += 1
                continue
            
            if token.type == TokenType.NOT:
                # Next token goes to must_not
                if i + 1 < len(tokens):
                    must_not_clauses.append(
                        self._token_to_query(tokens[i + 1], fields)
                    )
                    i += 2
                    continue
            
            if token.type == TokenType.LPAREN:
                # Find matching paren and recurse
                paren_depth = 1
                j = i + 1
                while j < len(tokens) and paren_depth > 0:
                    if tokens[j].type == TokenType.LPAREN:
                        paren_depth += 1
                    elif tokens[j].type == TokenType.RPAREN:
                        paren_depth -= 1
                    j += 1
                
                # Unbalanced: auto-close the group at the end of the query
                end = j - 1 if paren_depth == 0 else j
                sub_tokens = tokens[i + 1:end]
                sub_query = self._parse_expression(sub_tokens, fields)
                
                if current_operator == TokenType.OR:
                    should_clauses.append(sub_query)
                else:
                    must_clauses.append(sub_query)
                
                i = j
                continue
            
            if token.type == TokenType.PROXIMITY:
                # Proximity: previous term /N next term
                if must_clauses and i + 1 < len(tokens):
                    prev_query = must_clauses.pop()
                    next_query = self._token_to_query(tokens[i + 1], fields)
                    
                    proximity = self._parse_proximity(token.value)
                    span_query = self._build_span_near(
                        prev_query, next_query, proximity, fields[0]
                    )
                    must_clauses.append(span_query)
                    i += 2
                    continue
            
            # Regular term/phrase
            term_query = self._token_to_query(token, fields)
            
            if current_operator == TokenType.OR:
                should_clauses.append(term_query)
            else:
                must_clauses.append(term_query)
            
            i += 1
        
        # Build final bool query
        bool_query = {"bool": {}}
        
        if must_clauses:
            bool_query["bool"]["must"] = must_clauses
        if should_clauses:
            bool_query["bool"]["should"] = should_clauses
            bool_query["bool"]["minimum_should_match"] = 1
        if must_not_clauses:
            bool_query["bool"]["must_not"] = must_not_clauses
        
        return bool_query
    
    def _token_to_query(self, token: Token, fields: List[str]) -> dict:
        """Convert a single token to ES query."""
        
        if token.type == TokenType.TERM:
            return {
                "multi_match": {
                    "query": token.value,
                    "fields": fields,
                    "type": "best_fields"
                }
            }
        
        if token.type == TokenType.PHRASE:
            return {
                "multi_match": {
                    "query": token.value,
                    "fields": fields,
                    "type": "phrase"
                }
            }
        
        if token.type == TokenType.FIELD:
            # Parse field:value
            field, value = token.value.split(':', 1)
            return {
                "match": {field: value}
            }
        
        if token.type == TokenType.CITATION:
            # Handle citation search
            prefix, citation = token.value.split(':', 1)
            if prefix.lower() == 'cites':
                return {
                    "term": {"citations.cited_id": citation}
                }
            elif prefix.lower() == 'citedby':
                return {
                    "term": {"citations.citing_id": citation}
                }
        
        if token.type == TokenType.RANGE:
            # Parse [2020 TO 2024]
            match = re.match(r'\[(\S+)\s+TO\s+(\S+)\]', token.value)
            if match:
                return {
                    "range": {
                        "date": {
                            "gte": match.group(1),
                            "lte": match.group(2)
                        }
                    }
                }
        
        return {"match_all": {}}
    
    def _parse_proximity(self, value: str) -> int:
        """Parse proximity value to integer."""
        if value == 's':
            return 15  # Same sentence ≈ 15 words
        if value == 'p':
            return 50  # Same paragraph ≈ 50 words
        return int(value)
    
    def _build_span_near(
        self,
        query1: dict,
        query2: dict,
        slop: int,
        field: str
    ) -> dict:
        """Build span_near query for proximity search."""
        
        # Extract terms from queries
        term1 = self._extract_term(query1)
        term2 = self._extract_term(query2)
        
        return {
            "span_near": {
                "clauses": [
                    {"span_term": {field: term1}},
                    {"span_term": {field: term2}}
                ],
                "slop": slop,
                "in_order": False
            }
        }
    
    def _extract_term(self, query: dict) -> str:
        """Extract term from a query structure."""
        if "multi_match" in query:
            return query["multi_match"]["query"]
        if "match" in query:
            return list(query["match"].values())[0]
        return ""


# =============================================================================
# Usage Example
# =============================================================================

async def search_legal_documents(
    query_string: str,
    tenant_id: str,
    filters: dict = None
) -> dict:
    """
    Execute a legal search query.
    
    Example queries:
    - "(negligence /10 liability) AND damages"
    - "breach AND contract NOT employment"
    - "cites:410 U.S. 113"
    - "judge:Ginsburg AND civil rights"
    """
    
    builder = LegalQueryBuilder()
    
    es_query = builder.build_query(
        query_string,
        tenant_id=tenant_id,
        search_fields=["content", "title^3", "summary^2", "headnotes^2"]
    )
    
    # Add additional filters
    if filters:
        if filters.get("jurisdiction"):
            es_query["query"]["bool"]["filter"].append({
                "term": {"jurisdiction": filters["jurisdiction"]}
            })
        
        if filters.get("date_from") or filters.get("date_to"):
            date_range = {}
            if filters.get("date_from"):
                date_range["gte"] = filters["date_from"]
            if filters.get("date_to"):
                date_range["lte"] = filters["date_to"]
            es_query["query"]["bool"]["filter"].append({
                "range": {"decision_date": date_range}
            })
        
        if filters.get("document_type"):
            es_query["query"]["bool"]["filter"].append({
                "term": {"doc_type": filters["document_type"]}
            })
    
    # Add highlighting for snippets
    es_query["highlight"] = {
        "fields": {
            "content": {
                "fragment_size": 200,
                "number_of_fragments": 3
            },
            "headnotes": {}
        },
        "pre_tags": ["<mark>"],
        "post_tags": ["</mark>"]
    }
    
    return es_query

# security/tenant_isolation.py

"""
Multi-tenant isolation for legal document search.

Applies: Week 7, Day 1 - Document modeling (a separate index per tenant)
         Week 7, Day 3 - Filter context (the mandatory tenant filter)
         Tenant security itself is new in this capstone (Week 9 goes deeper)
"""

from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Optional, List
from functools import wraps
import logging

logger = logging.getLogger(__name__)


class AuthenticationError(Exception):
    """Missing or invalid token."""


class SecurityError(Exception):
    """Tenant isolation violation."""


@dataclass(frozen=True)
class TenantContext:
    """Immutable tenant context from JWT."""
    tenant_id: str
    user_id: str
    roles: List[str]
    
    @property
    def private_index(self) -> str:
        """Get tenant-specific private index name."""
        return f"private_docs_{self.tenant_id}"
    
    @property
    def can_search_public(self) -> bool:
        """Check if user can search public documents."""
        return "search:public" in self.roles
    
    @property
    def can_search_private(self) -> bool:
        """Check if user can search private documents."""
        return "search:private" in self.roles


class TenantIsolationMiddleware:
    """
    Middleware to enforce tenant isolation.
    
    This runs BEFORE any search service code.
    """
    
    def __init__(self, jwt_validator):
        self.jwt_validator = jwt_validator
    
    async def __call__(self, request):
        """Extract and validate tenant context."""
        
        # Get token from header
        auth_header = request.headers.get("Authorization", "")
        if not auth_header.startswith("Bearer "):
            raise AuthenticationError("Missing token")
        
        token = auth_header[7:]
        
        # Validate and extract claims
        claims = await self.jwt_validator.validate(token)
        
        # Create immutable context
        context = TenantContext(
            tenant_id=claims["tenant_id"],
            user_id=claims["user_id"],
            roles=claims.get("roles", [])
        )
        
        # Attach to request (cannot be modified)
        request.state.tenant = context
        
        # Log for audit
        logger.info(
            "Request authenticated",
            extra={
                "tenant_id": context.tenant_id,
                "user_id": context.user_id,
                "path": request.url.path
            }
        )
        
        return context


def require_tenant_isolation(func):
    """
    Decorator ensuring tenant isolation in search functions.
    """
    
    @wraps(func)
    async def wrapper(*args, **kwargs):
        # Get tenant context (must exist)
        tenant: TenantContext = kwargs.get("tenant")
        if not tenant:
            raise SecurityError("Tenant context required")
        
        # Verify tenant_id is not overridable via the query or filters
        query = {
            "q": kwargs.get("query_string", ""),
            "f": kwargs.get("filters") or {},
        }
        if _contains_tenant_override(query):
            logger.warning(
                "Blocked tenant override attempt",
                extra={
                    "tenant_id": tenant.tenant_id,
                    "attempted_query": query
                }
            )
            raise SecurityError("Tenant override not allowed")
        
        return await func(*args, **kwargs)
    
    return wrapper


def _contains_tenant_override(query: dict) -> bool:
    """Check if query tries to override tenant filter."""
    
    query_str = str(query).lower()
    
    # Block any attempt to query tenant_id field directly
    if "tenant_id" in query_str:
        return True
    
    # Block attempts to query other tenant indices
    if "private_docs_" in query_str:
        return True
    
    return False


class IsolatedSearchService:
    """
    Search service with built-in tenant isolation.
    """
    
    def __init__(self, es_client):
        self.es = es_client
        self.public_index = "public_legal_docs"
    
    @require_tenant_isolation
    async def search(
        self,
        query_string: str,
        tenant: TenantContext,
        include_public: bool = True,
        include_private: bool = True,
        filters: dict = None
    ) -> dict:
        """
        Execute search with automatic tenant isolation.
        
        The tenant context comes from JWT, not from the request body.
        It cannot be overridden by the caller.
        """
        
        indices = []
        
        # Public documents (if permitted and requested)
        if include_public and tenant.can_search_public:
            indices.append(self.public_index)
        
        # Private documents (tenant-specific index)
        if include_private and tenant.can_search_private:
            indices.append(tenant.private_index)
        
        if not indices:
            return {"hits": [], "total": 0}
        
        # Build query with MANDATORY tenant filter
        es_query = self._build_isolated_query(
            query_string,
            tenant,
            filters
        )
        
        # Execute search
        response = await self.es.search(
            index=",".join(indices),
            body=es_query
        )
        
        # Audit log
        await self._audit_log(
            tenant=tenant,
            query=query_string,
            indices=indices,
            result_count=response["hits"]["total"]["value"]
        )
        
        return self._process_response(response)
    
    def _build_isolated_query(
        self,
        query_string: str,
        tenant: TenantContext,
        filters: dict
    ) -> dict:
        """Build query with mandatory tenant isolation."""
        
        # Parse the user's query
        parsed_query = LegalQueryBuilder().build_query(
            query_string,
            tenant_id=None,  # We add filter separately
            search_fields=["content", "title^3", "summary^2"]
        )
        
        # CRITICAL: Add tenant filter that CANNOT be bypassed
        # This filter is added by the service, not from user input
        tenant_filter = {
            "bool": {
                "should": [
                    # Public documents (no tenant restriction)
                    {"term": {"is_public": True}},
                    # Private documents (must match tenant)
                    {
                        "bool": {
                            "must": [
                                {"term": {"is_public": False}},
                                {"term": {"tenant_id": tenant.tenant_id}}
                            ]
                        }
                    }
                ],
                "minimum_should_match": 1
            }
        }
        
        # Inject filter
        parsed_query["query"]["bool"]["filter"].append(tenant_filter)
        
        # Add any additional filters
        if filters:
            for key, value in filters.items():
                if key in ["jurisdiction", "doc_type", "court"]:
                    parsed_query["query"]["bool"]["filter"].append({
                        "term": {key: value}
                    })
        
        return parsed_query
    
    async def _audit_log(
        self,
        tenant: TenantContext,
        query: str,
        indices: List[str],
        result_count: int
    ):
        """Log search for audit trail."""
        
        await self.es.index(
            index="search_audit_log",
            document={
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "tenant_id": tenant.tenant_id,
                "user_id": tenant.user_id,
                "query": query,
                "indices_searched": indices,
                "result_count": result_count,
                "client_ip": "extracted_from_request"
            }
        )
    
    def _process_response(self, response: dict) -> dict:
        """Process ES response, ensuring no cross-tenant leakage."""
        
        hits = []
        for hit in response["hits"]["hits"]:
            # Double-check: never return tenant_id in response
            source = hit["_source"]
            if "tenant_id" in source:
                del source["tenant_id"]
            
            hits.append({
                "id": hit["_id"],
                "score": hit["_score"],
                "source": source,
                "highlights": hit.get("highlight", {})
            })
        
        return {
            "hits": hits,
            "total": response["hits"]["total"]["value"]
        }

# tests/test_tenant_isolation.py

"""
Critical security tests for tenant isolation.

These tests MUST pass before any deployment.
"""

import pytest


class TestTenantIsolation:
    """Tests that verify tenant isolation cannot be bypassed."""
    
    async def test_cannot_search_other_tenant_documents(
        self,
        search_service,
        tenant_a,
        tenant_b
    ):
        """Tenant A cannot find Tenant B's documents."""
        
        # Tenant B uploads a document
        doc = await upload_document(
            tenant=tenant_b,
            content="This is a secret merger agreement"
        )
        
        # Tenant A searches for it
        results = await search_service.search(
            query_string="secret merger agreement",
            tenant=tenant_a,
            include_private=True
        )
        
        # Must NOT find it
        assert doc["id"] not in [h["id"] for h in results["hits"]]
    
    async def test_cannot_access_by_id_guess(
        self,
        document_service,
        tenant_a,
        tenant_b
    ):
        """Tenant A cannot access Tenant B's document by ID."""
        
        # Tenant B's document
        doc = await upload_document(tenant=tenant_b, content="Secret")
        
        # Tenant A tries to access by ID
        with pytest.raises(NotFoundError):
            await document_service.get_document(
                document_id=doc["id"],
                tenant=tenant_a
            )
    
    async def test_cannot_override_tenant_filter(
        self,
        search_service,
        tenant_a
    ):
        """Query cannot override tenant filter."""
        
        # Try to inject tenant_id in query
        with pytest.raises(SecurityError):
            await search_service.search(
                query_string="contract",
                tenant=tenant_a,
                filters={"tenant_id": "other_tenant"}  # Should be blocked
            )
    
    async def test_cannot_query_other_tenant_index(
        self,
        search_service,
        tenant_a
    ):
        """Cannot query another tenant's index directly."""
        
        # Try to search other tenant's index
        with pytest.raises(SecurityError):
            await search_service.search_raw(
                index="private_docs_other_tenant",
                query={"match_all": {}},
                tenant=tenant_a
            )
    
    async def test_public_documents_visible_to_all(
        self,
        search_service,
        tenant_a,
        tenant_b
    ):
        """Public documents are searchable by all tenants."""
        
        # Search public case law
        results_a = await search_service.search(
            query_string="Roe v. Wade",
            tenant=tenant_a,
            include_public=True,
            include_private=False  # Compare public results only
        )
        
        results_b = await search_service.search(
            query_string="Roe v. Wade",
            tenant=tenant_b,
            include_public=True,
            include_private=False
        )
        
        # Both should find the same public documents
        assert results_a["total"] > 0
        assert results_a["total"] == results_b["total"]

# citations/citation_service.py

"""
Citation network search and analysis.

Applies: Week 7, Day 2 - Indexing Pipeline (citation extraction)
         Week 7, Day 4 - Aggregations, as in facets (treatment breakdowns)
         Graph-like queries over citation edges are new in this capstone
"""

from dataclasses import dataclass
from typing import List, Optional, Dict
from enum import Enum


class CitationTreatment(Enum):
    """How a case was cited."""
    FOLLOWED = "followed"
    DISTINGUISHED = "distinguished"
    OVERRULED = "overruled"
    QUESTIONED = "questioned"
    MENTIONED = "mentioned"


@dataclass
class Citation:
    """A citation between two legal documents."""
    citing_id: str
    citing_name: str
    cited_id: str
    cited_name: str
    treatment: CitationTreatment
    context: str  # Surrounding text
    date: str


@dataclass
class CitationAnalysis:
    """Analysis of a document's citation profile."""
    document_id: str
    document_name: str
    total_citations: int
    treatment_breakdown: Dict[str, int]
    is_good_law: bool
    overruled_by: Optional[str]
    key_citing_cases: List[dict]


class CitationService:
    """
    Manages citation network queries.
    """
    
    def __init__(self, es_client):
        self.es = es_client
        self.docs_index = "legal_documents"
        self.citations_index = "citation_edges"
    
    async def find_citing_cases(
        self,
        document_id: str,
        treatment: Optional[CitationTreatment] = None,
        limit: int = 100
    ) -> List[dict]:
        """
        Find all cases that cite a given document.
        
        This is the "cited by" query - who cites this case?
        """
        
        query = {
            "query": {
                "bool": {
                    "must": [
                        {"term": {"cited_id": document_id}}
                    ]
                }
            },
            "sort": [
                {"citation_date": "desc"}
            ],
            "size": limit
        }
        
        if treatment:
            query["query"]["bool"]["must"].append({
                "term": {"treatment": treatment.value}
            })
        
        response = await self.es.search(
            index=self.citations_index,
            body=query
        )
        
        # Enrich with document details
        citing_ids = [hit["_source"]["citing_id"] for hit in response["hits"]["hits"]]
        
        docs = await self._get_documents(citing_ids)
        
        results = []
        for hit in response["hits"]["hits"]:
            citation = hit["_source"]
            doc = docs.get(citation["citing_id"], {})
            
            results.append({
                "document_id": citation["citing_id"],
                "name": doc.get("name", citation["citing_id"]),
                "court": doc.get("court"),
                "date": doc.get("decision_date"),
                "treatment": citation["treatment"],
                "context": citation.get("context_snippet", "")
            })
        
        return results
    
    async def find_cited_cases(
        self,
        document_id: str,
        limit: int = 100
    ) -> List[dict]:
        """
        Find all cases cited by a given document.
        
        This is the "cites" query - what does this case cite?
        """
        
        query = {
            "query": {
                "term": {"citing_id": document_id}
            },
            "size": limit
        }
        
        response = await self.es.search(
            index=self.citations_index,
            body=query
        )
        
        # Enrich with document details
        cited_ids = [hit["_source"]["cited_id"] for hit in response["hits"]["hits"]]
        
        docs = await self._get_documents(cited_ids)
        
        results = []
        for hit in response["hits"]["hits"]:
            citation = hit["_source"]
            doc = docs.get(citation["cited_id"], {})
            
            results.append({
                "document_id": citation["cited_id"],
                "name": doc.get("name", citation["cited_id"]),
                "court": doc.get("court"),
                "date": doc.get("decision_date"),
                "treatment": citation["treatment"]
            })
        
        return results
    
    async def analyze_citation_profile(
        self,
        document_id: str
    ) -> CitationAnalysis:
        """
        Analyze the citation profile of a document.
        
        Determines if case is "good law" and how it's been treated.
        """
        
        # Get the document
        doc = await self._get_document(document_id)
        
        # Get citation treatment aggregation
        agg_query = {
            "query": {
                "term": {"cited_id": document_id}
            },
            "size": 0,
            "track_total_hits": True,  # exact count; the default caps at 10,000
            "aggs": {
                "treatments": {
                    "terms": {"field": "treatment"}
                },
                "recent_overruled": {
                    "filter": {"term": {"treatment": "overruled"}},
                    "aggs": {
                        "cases": {
                            "top_hits": {
                                "size": 1,
                                "sort": [{"citation_date": "desc"}]
                            }
                        }
                    }
                },
                "key_citations": {
                    "top_hits": {
                        "size": 10,
                        "sort": [{"citing_court_rank": "desc"}]
                    }
                }
            }
        }
        
        response = await self.es.search(
            index=self.citations_index,
            body=agg_query
        )
        
        aggs = response["aggregations"]
        
        # Build treatment breakdown
        treatment_breakdown = {
            bucket["key"]: bucket["doc_count"]
            for bucket in aggs["treatments"]["buckets"]
        }
        
        # Determine if good law
        overruled_hits = aggs["recent_overruled"]["cases"]["hits"]["hits"]
        is_overruled = len(overruled_hits) > 0
        overruled_by = None
        
        if is_overruled:
            overruled_by = overruled_hits[0]["_source"]["citing_id"]
        
        # Key citing cases
        key_cases = [
            {
                "id": hit["_source"]["citing_id"],
                "treatment": hit["_source"]["treatment"]
            }
            for hit in aggs["key_citations"]["hits"]["hits"]
        ]
        
        return CitationAnalysis(
            document_id=document_id,
            document_name=doc.get("name", ""),
            total_citations=response["hits"]["total"]["value"],
            treatment_breakdown=treatment_breakdown,
            is_good_law=not is_overruled,
            overruled_by=overruled_by,
            key_citing_cases=key_cases
        )
    
    async def find_citation_chain(
        self,
        document_id: str,
        depth: int = 2,
        limit_per_level: int = 10
    ) -> dict:
        """
        Find citation chain up to N levels deep.
        
        Example: Cases citing cases that cite the target case.
        """
        
        chain = {
            "root": document_id,
            "levels": []
        }
        
        current_ids = [document_id]
        
        for level in range(depth):
            # Find all cases citing the current level
            query = {
                "query": {
                    "terms": {"cited_id": current_ids}
                },
                "aggs": {
                    "citing_cases": {
                        "terms": {
                            "field": "citing_id",
                            "size": limit_per_level
                        }
                    }
                },
                "size": 0
            }
            
            response = await self.es.search(
                index=self.citations_index,
                body=query
            )
            
            citing_ids = [
                bucket["key"]
                for bucket in response["aggregations"]["citing_cases"]["buckets"]
            ]
            
            if not citing_ids:
                break
            
            chain["levels"].append({
                "depth": level + 1,
                "count": len(citing_ids),
                "sample_ids": citing_ids[:5]
            })
            
            current_ids = citing_ids
        
        return chain
    
    async def _get_document(self, document_id: str) -> dict:
        """Get a single document by ID."""
        try:
            response = await self.es.get(
                index=self.docs_index,
                id=document_id
            )
            return response["_source"]
        except:
            return {}
    
    async def _get_documents(self, document_ids: List[str]) -> Dict[str, dict]:
        """Get multiple documents by ID."""
        if not document_ids:
            return {}
        
        response = await self.es.mget(
            index=self.docs_index,
            body={"ids": document_ids}
        )
        
        return {
            doc["_id"]: doc["_source"]
            for doc in response["docs"]
            if doc.get("found")
        }


# =============================================================================
# Citation Extraction Pipeline
# =============================================================================

class CitationExtractor:
    """
    Extracts citations from legal document text.
    
    Applies: Week 7, Day 2 - Indexing Pipeline
    """
    
    # Patterns for different citation formats
    PATTERNS = {
        # US Supreme Court: 410 U.S. 113
        "us_reports": r"(\d+)\s+U\.S\.\s+(\d+)",
        
        # Federal Reporter: 123 F.2d 456
        "federal_reporter": r"(\d+)\s+F\.(2d|3d)?\s+(\d+)",
        
        # State reporters: 123 Cal.App.4th 456
        "state_reporter": r"(\d+)\s+([A-Z][a-z]+\.?\s*(?:App\.?)?\s*(?:\d+[a-z]+)?)\s+(\d+)",
        
        # Case names: Brown v. Board of Education
        "case_name": r"([A-Z][a-z]+)\s+v\.\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)",
    }
    
    def extract_citations(self, text: str) -> List[dict]:
        """Extract all citations from document text."""
        
        import re
        
        citations = []
        
        for pattern_name, pattern in self.PATTERNS.items():
            for match in re.finditer(pattern, text):
                # Get surrounding context
                start = max(0, match.start() - 100)
                end = min(len(text), match.end() + 100)
                context = text[start:end]
                
                # Detect treatment from context
                treatment = self._detect_treatment(context)
                
                citations.append({
                    "raw_citation": match.group(0),
                    "pattern_type": pattern_name,
                    "context": context,
                    "treatment": treatment,
                    "position": match.start()
                })
        
        return citations
    
    def _detect_treatment(self, context: str) -> str:
        """Detect how a citation is being used."""
        
        context_lower = context.lower()
        
        # Negative treatments
        if any(word in context_lower for word in 
               ["overruled", "overrule", "rejected", "abrogated"]):
            return "overruled"
        
        if any(word in context_lower for word in 
               ["questioned", "doubted", "criticized"]):
            return "questioned"
        
        # Neutral/distinguishing
        if any(word in context_lower for word in 
               ["distinguished", "distinguishing", "unlike"]):
            return "distinguished"
        
        # Positive
        if any(word in context_lower for word in 
               ["followed", "following", "accord", "see also", "affirmed"]):
            return "followed"
        
        # Default
        return "mentioned"

# ingestion/document_processor.py

"""
Document processing pipeline for legal documents.

Applies: Week 7, Day 2 - Indexing Pipeline
"""

from dataclasses import dataclass
from typing import Optional, List
from enum import Enum
import asyncio
import logging

logger = logging.getLogger(__name__)


class DocumentType(Enum):
    PDF_TEXT = "pdf_text"
    PDF_SCANNED = "pdf_scanned"
    WORD = "word"
    IMAGE = "image"
    UNKNOWN = "unknown"


@dataclass
class ProcessedDocument:
    """Result of document processing."""
    document_id: str
    tenant_id: str
    original_filename: str
    content: str
    metadata: dict
    citations: List[dict]
    page_count: int
    word_count: int
    ocr_confidence: Optional[float]
    processing_time_ms: int


class DocumentProcessor:
    """
    Processes uploaded documents for indexing.
    """
    
    def __init__(
        self,
        s3_client,
        tika_client,
        ocr_service,
        citation_extractor,
        kafka_producer
    ):
        self.s3 = s3_client
        self.tika = tika_client
        self.ocr = ocr_service
        self.citations = citation_extractor
        self.kafka = kafka_producer
    
    async def process_document(
        self,
        document_id: str,
        tenant_id: str,
        s3_key: str,
        filename: str
    ) -> ProcessedDocument:
        """
        Process a document through the full pipeline.
        """
        
        import time
        start_time = time.time()
        
        # Download from S3
        file_bytes = await self.s3.download(s3_key)
        
        # Determine document type
        doc_type = self._detect_document_type(filename, file_bytes)
        
        logger.info(
            f"Processing document {document_id}",
            extra={"type": doc_type.value, "size": len(file_bytes)}
        )
        
        # Extract content based on type
        if doc_type == DocumentType.PDF_TEXT:
            content, metadata = await self._extract_pdf_text(file_bytes)
            ocr_confidence = None
            
        elif doc_type == DocumentType.PDF_SCANNED:
            content, metadata, ocr_confidence = await self._process_scanned_pdf(
                file_bytes
            )
            
        elif doc_type == DocumentType.WORD:
            content, metadata = await self._extract_word(file_bytes)
            ocr_confidence = None
            
        elif doc_type == DocumentType.IMAGE:
            content, ocr_confidence = await self._process_image(file_bytes)
            metadata = {}
            
        else:
            # Fallback to Tika
            content, metadata = await self._extract_with_tika(file_bytes)
            ocr_confidence = None
        
        # Extract citations
        citations = self.citations.extract_citations(content)
        
        # Build processed document
        processing_time = int((time.time() - start_time) * 1000)
        
        processed = ProcessedDocument(
            document_id=document_id,
            tenant_id=tenant_id,
            original_filename=filename,
            content=content,
            metadata=metadata,
            citations=citations,
            page_count=metadata.get("page_count", 1),
            word_count=len(content.split()),
            ocr_confidence=ocr_confidence,
            processing_time_ms=processing_time
        )
        
        # Publish to Kafka for indexing
        await self._publish_for_indexing(processed)
        
        return processed
    
    def _detect_document_type(
        self,
        filename: str,
        file_bytes: bytes
    ) -> DocumentType:
        """Detect the type of document for processing."""
        
        ext = filename.lower().split(".")[-1]
        
        if ext in ("doc", "docx"):
            return DocumentType.WORD
        
        if ext in ("jpg", "jpeg", "png", "tiff", "bmp"):
            return DocumentType.IMAGE
        
        if ext == "pdf":
            # Check if PDF has extractable text
            if self._pdf_has_text(file_bytes):
                return DocumentType.PDF_TEXT
            else:
                return DocumentType.PDF_SCANNED
        
        return DocumentType.UNKNOWN
    
    def _pdf_has_text(self, file_bytes: bytes) -> bool:
        """Check if PDF contains extractable text or is scanned."""
        
        import fitz  # PyMuPDF
        
        doc = fitz.open(stream=file_bytes, filetype="pdf")
        
        # Check first few pages
        for page_num in range(min(3, len(doc))):
            page = doc[page_num]
            text = page.get_text()
            if len(text.strip()) > 100:
                return True
        
        return False
    
    async def _extract_pdf_text(self, file_bytes: bytes) -> tuple:
        """Extract text from text-based PDF."""
        
        import fitz
        
        doc = fitz.open(stream=file_bytes, filetype="pdf")
        
        content_parts = []
        for page in doc:
            content_parts.append(page.get_text())
        
        content = "\n\n".join(content_parts)
        
        metadata = {
            "page_count": len(doc),
            "title": doc.metadata.get("title", ""),
            "author": doc.metadata.get("author", ""),
            "creation_date": doc.metadata.get("creationDate", "")
        }
        
        return content, metadata
    
    async def _process_scanned_pdf(self, file_bytes: bytes) -> tuple:
        """Process scanned PDF through OCR."""
        
        import fitz
        
        doc = fitz.open(stream=file_bytes, filetype="pdf")
        
        content_parts = []
        confidences = []
        
        for page_num, page in enumerate(doc):
            # Extract page as image
            pix = page.get_pixmap(dpi=300)
            image_bytes = pix.tobytes("png")
            
            # OCR the image
            text, confidence = await self.ocr.process_image(image_bytes)
            
            content_parts.append(text)
            confidences.append(confidence)
        
        content = "\n\n".join(content_parts)
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0
        
        metadata = {"page_count": len(doc)}
        
        return content, metadata, avg_confidence
    
    async def _process_image(self, file_bytes: bytes) -> tuple:
        """Process single image through OCR."""
        
        text, confidence = await self.ocr.process_image(file_bytes)
        return text, confidence
    
    async def _extract_word(self, file_bytes: bytes) -> tuple:
        """Extract content from Word document."""
        
        result = await self.tika.extract(file_bytes, "application/msword")
        
        return result["content"], result.get("metadata", {})
    
    async def _extract_with_tika(self, file_bytes: bytes) -> tuple:
        """Fallback extraction with Apache Tika."""
        
        result = await self.tika.extract(file_bytes)
        return result.get("content", ""), result.get("metadata", {})
    
    async def _publish_for_indexing(self, document: ProcessedDocument):
        """Publish processed document to Kafka for indexing."""
        
        event = {
            "event_type": "document.created",
            "document_id": document.document_id,
            "tenant_id": document.tenant_id,
            "content": document.content,
            "metadata": document.metadata,
            "citations": document.citations,
            "stats": {
                "page_count": document.page_count,
                "word_count": document.word_count,
                "ocr_confidence": document.ocr_confidence,
                "processing_time_ms": document.processing_time_ms
            }
        }
        
        await self.kafka.produce(
            topic="document.created",
            key=document.document_id,
            value=event
        )