components/text_splitter.py

"""
text_splitter.py
----------------
Splits LangChain Document objects into smaller, overlapping chunks.
Enhanced with structured metadata, validation, and better organization.

Chunk size and overlap are driven by config.py to keep the logic
configurable without touching source code.

Topic Detection (NEW):
- Automatically detects section headings and keywords
- Tags chunks with topic metadata (e.g., "topic: culture", "topic: economy")
- Enables more precise retrieval by topic matching
"""

import logging
import re
from typing import List

from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter

from app.config import CHUNK_SIZE, CHUNK_OVERLAP

logger = logging.getLogger(__name__)


def _validate_chunk(chunk: Document, min_length: int = 10) -> bool:
    """
    Validate that a chunk is meaningful.
    
    Args:
        chunk: Document chunk to validate.
        min_length: Minimum character count (default 10).
    
    Returns:
        True if chunk passes validation, False otherwise.
    """
    # Check length
    if len(chunk.page_content) < min_length:
        return False
    
    # Check that it's not just whitespace or punctuation
    meaningful_chars = sum(
        1 for c in chunk.page_content 
        if c.isalnum()
    )
    if meaningful_chars < min_length:
        return False
    
    return True


def _ends_with_complete_sentence(text: str) -> bool:
    """
    Check if text ends with a complete sentence (ends with sentence terminator).
    
    Args:
        text: Text to check.
    
    Returns:
        True if text ends with `.`, `!`, `?`, or other sentence terminators.
    """
    text = text.rstrip()
    sentence_terminators = {'.', '!', '?', ':', ';'}
    return len(text) > 0 and text[-1] in sentence_terminators


def _truncate_at_sentence_boundary(text: str) -> str:
    """
    Truncate text at the last complete sentence boundary.
    
    If text ends mid-sentence, find the last sentence terminator and truncate there.
    Falls back to original text if no boundary found.
    
    Args:
        text: Text potentially ending mid-sentence.
    
    Returns:
        Text truncated at a sentence boundary, or original text if no boundary found.
    """
    sentence_terminators = {'.', '!', '?'}
    
    # Look backwards for the last sentence terminator
    for i in range(len(text) - 1, -1, -1):
        if text[i] in sentence_terminators:
            # Return text up to and including the terminator
            return text[:i + 1].rstrip()
    
    # No sentence terminator found; return original text
    return text


def _fix_chunk_boundaries(chunk: Document) -> Document | None:
    """
    Ensure a chunk ends at a sentence boundary.
    
    If chunk ends mid-sentence, truncate at the last complete sentence.
    If truncation leaves too little content, return None (discard chunk).
    
    Args:
        chunk: Document chunk to validate/fix.
    
    Returns:
        Fixed chunk with complete sentences, or None if chunk becomes too small.
    """
    content = chunk.page_content
    min_length = 50  # Minimum characters after boundary adjustment
    
    # If already ends at sentence boundary, no fix needed
    if _ends_with_complete_sentence(content):
        return chunk
    
    # Try to fix by truncating at sentence boundary
    fixed_content = _truncate_at_sentence_boundary(content)
    
    # Validate the fixed content has enough length
    if len(fixed_content) < min_length:
        logger.debug(
            "Chunk too short after sentence boundary fix (%d chars, min %d)",
            len(fixed_content),
            min_length,
        )
        return None
    
    # Update chunk with fixed content
    chunk.page_content = fixed_content
    return chunk


def _detect_topic(text: str) -> str | None:
    """
    Detect the topic/section of a chunk by looking for section headers
    and common topic keywords.
    
    Args:
        text: Chunk content to analyze.
    
    Returns:
        Detected topic string (e.g., "culture", "economy") or None.
    """
    # Look for section headers (lines that look like headings)
    # Match patterns like "### Culture" or "## ECONOMY" or "Culture:"
    header_patterns = [
        r"^#+\s+([A-Za-z\s]+?)(?:\n|$)",  # Markdown headers
        r"^([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)\s*:\s*$",  # "Culture: " at line start
        r"^([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)\s*$",  # "Culture" alone on a line
    ]
    
    for pattern in header_patterns:
        match = re.search(pattern, text, re.MULTILINE)
        if match:
            topic = match.group(1).strip().lower()
            if len(topic) > 2:  # Filter out single letters
                return topic
    
    # Fallback: look for common topic keywords in text
    topic_keywords = {
        "culture": ["culture", "cultural", "tradition", "custom", "heritage"],
        "economy": ["economy", "economic", "trade", "commerce", "market", "gdp"],
        "geography": ["geography", "geographical", "location", "region", "area"],
        "history": ["history", "historical", "past", "century"],
        "population": ["population", "demographic", "resident", "inhabitant"],
        "government": ["government", "political", "administration", "state", "federal"],
        "religion": ["religion", "religious", "faith", "belief"],
        "education": ["education", "school", "university", "college"],
        "climate": ["climate", "weather", "temperature", "precipitation"],
        "language": ["language", "linguistic", "speak", "dialect"],
    }
    
    text_lower = text.lower()
    for topic, keywords in topic_keywords.items():
        # Count keyword occurrences
        matches = sum(1 for kw in keywords if kw in text_lower)
        if matches >= 2:  # If 2+ keywords match, assign this topic
            return topic
    
    return None



def _enrich_chunk_metadata(
    chunk: Document, 
    chunk_index: int,
    total_chunks: int,
) -> Document:
    """
    Add structured metadata to a chunk for retrieval and attribution.
    
    Args:
        chunk: Document chunk.
        chunk_index: Position of this chunk in sequence.
        total_chunks: Total number of chunks from same source.
    
    Returns:
        Chunk with enriched metadata.
    """
    chunk.metadata["chunk_id"] = chunk_index
    chunk.metadata["chunk_total"] = total_chunks
    
    # Add length info for relevance scoring
    chunk.metadata["chunk_chars"] = len(chunk.page_content)
    chunk.metadata["chunk_words"] = len(chunk.page_content.split())
    
    # Detect and tag topic
    topic = _detect_topic(chunk.page_content)
    if topic:
        chunk.metadata["topic"] = topic
        logger.debug(f"Detected topic '{topic}' in chunk {chunk_index}")
    
    # Preview for debugging (first 50 chars)
    preview = chunk.page_content[:50].replace('\n', ' ')
    chunk.metadata["chunk_preview"] = preview
    
    return chunk


def split_documents(
    documents: List[Document],
    chunk_size: int = CHUNK_SIZE,
    chunk_overlap: int = CHUNK_OVERLAP,
) -> List[Document]:
    """
    Split a list of Documents into smaller chunks with overlap.

    Each output chunk:
    - Inherits metadata of its parent document (source, page number, etc.)
    - Receives structured chunk-level metadata (chunk_id, preview, etc.)
    - Is validated for meaningfulness
    - Preserves source attribution for retrieval

    Args:
        documents:    List of LangChain Document objects to split.
        chunk_size:   Maximum number of characters per chunk.
        chunk_overlap: Number of characters shared between consecutive chunks.

    Returns:
        List of validated, enriched chunked Document objects.
    """
    if not documents:
        logger.warning("split_documents called with an empty document list.")
        return []

    # Use intelligent separator hierarchy for clean chunk boundaries
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        # Prefer splitting at semantic boundaries:
        # 1. Paragraph breaks (double newline)
        # 2. Single newlines
        # 3. Sentence endings
        # 4. Word boundaries
        # 5. Last resort: character level
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=len,
        add_start_index=True,  # Track char position in source
    )

    # Perform initial split
    raw_chunks = splitter.split_documents(documents)
    
    # Filter and enrich chunks
    valid_chunks = []
    
    for chunk in raw_chunks:
        # Validate chunk
        if not _validate_chunk(chunk):
            logger.debug(
                "Skipping invalid chunk from '%s' (too short/empty)",
                chunk.metadata.get("source", "unknown"),
            )
            continue
        
        # Fix sentence boundaries: ensure chunk ends at complete sentence
        fixed_chunk = _fix_chunk_boundaries(chunk)
        if fixed_chunk is None:
            logger.debug(
                "Skipping chunk from '%s' (too short after sentence boundary fix)",
                chunk.metadata.get("source", "unknown"),
            )
            continue
        
        valid_chunks.append(fixed_chunk)
    
    # Add structured metadata to each chunk
    for source_doc_chunks in _group_chunks_by_source(valid_chunks):
        for chunk_idx, chunk in enumerate(source_doc_chunks):
            _enrich_chunk_metadata(chunk, chunk_idx, len(source_doc_chunks))
    
    logger.info(
        "Split %d document(s) → %d raw chunks → %d valid chunks "
        "(size=%d, overlap=%d, with sentence boundary validation)",
        len(documents),
        len(raw_chunks),
        len(valid_chunks),
        chunk_size,
        chunk_overlap,
    )
    
    return valid_chunks


def _group_chunks_by_source(chunks: List[Document]) -> List[List[Document]]:
    """
    Group chunks by their source document for per-source indexing.
    
    Args:
        chunks: List of chunks.
    
    Returns:
        List of chunk groups, each group from same source.
    """
    from collections import defaultdict
    
    groups = defaultdict(list)
    for chunk in chunks:
        source = chunk.metadata.get("source", "unknown")
        groups[source].append(chunk)
    
    return list(groups.values())