"""
Content-type-aware chunking service.

Uses semantic-text-splitter with HuggingFace tokenizers to produce
chunks that respect document structure boundaries per library type.
"""

import logging
from typing import Optional

from library.metrics import CHUNKS_CREATED_TOTAL, CHUNK_SIZE_TOKENS

from .parsers import ParseResult, TextBlock

logger = logging.getLogger(__name__)

# Default tokenizer when no model-specific tokenizer is available
DEFAULT_TOKENIZER = "bert-base-uncased"

# Boundary markers used to detect structural elements in text
_BOUNDARY_PATTERNS = {
    "chapter": [
        r"(?m)^chapter\s+\d+",
        r"(?m)^CHAPTER\s+\d+",
        r"(?m)^Chapter\s+\w+",
    ],
    "scene": [r"(?m)^\*\s*\*\s*\*", r"(?m)^---+$", r"(?m)^###"],
    "section": [
        r"(?m)^#{1,3}\s+",
        r"(?m)^\d+\.\d*\s+\w",
        r"(?m)^Section\s+\d+",
    ],
    "subsection": [r"(?m)^#{4,6}\s+", r"(?m)^\d+\.\d+\.\d+"],
    "entry": [
        r"(?m)^\d{4}-\d{2}-\d{2}",
        r"(?m)^(January|February|March|April|May|June|July|August|September|October|November|December)\s+\d",
    ],
    "song": [r"(?m)^Track\s+\d+", r"(?m)^\[.*?\]$"],
    "verse": [r"(?m)^\[Verse", r"(?m)^\[Chorus"],
}


class ChunkResult:
    """Result of chunking a document."""

    def __init__(
        self,
        chunks: list[str],
        chunk_page_map: dict[int, int],
        strategy: str,
    ):
        """
        :param chunks: List of chunk text strings.
        :param chunk_page_map: Mapping of chunk_index -> source page number.
        :param strategy: The chunking strategy used.
        """
        self.chunks = chunks
        self.chunk_page_map = chunk_page_map
        self.strategy = strategy

    def __len__(self) -> int:
        return len(self.chunks)


class ContentTypeChunker:
    """
    Content-type-aware document chunker.

    Dispatches to different chunking strategies based on the library's
    chunking configuration, using semantic-text-splitter for
    token-aware splitting.
    """

    def __init__(self, tokenizer_name: Optional[str] = None):
        """
        :param tokenizer_name: HuggingFace tokenizer name for token counting.
        """
        self._tokenizer_name = tokenizer_name or DEFAULT_TOKENIZER
        self._splitter_cache: dict[tuple[int, int], object] = {}

    def chunk(
        self,
        parse_result: ParseResult,
        chunking_config: dict,
        library_type: str = "",
    ) -> ChunkResult:
        """
        Chunk parsed document text using the library's chunking config.

        :param parse_result: ParseResult from the document parser.
        :param chunking_config: Library chunking configuration dict.
        :param library_type: Library type for metrics labeling.
        :returns: ChunkResult with chunk texts and page mapping.
        """
        strategy = chunking_config.get("strategy", "section_aware")
        chunk_size = chunking_config.get("chunk_size", 512)
        chunk_overlap = chunking_config.get("chunk_overlap", 64)

        # Combine all text blocks into a single document text,
        # tracking page boundaries for chunk-page mapping
        full_text, page_offsets = self._combine_text_blocks(parse_result.text_blocks)

        if not full_text.strip():
            logger.warning("No text to chunk strategy=%s", strategy)
            return ChunkResult(chunks=[], chunk_page_map={}, strategy=strategy)

        logger.info(
            "Chunking text strategy=%s chunk_size=%d overlap=%d total_chars=%d",
            strategy,
            chunk_size,
            chunk_overlap,
            len(full_text),
        )

        # Get or create the text splitter for this size/overlap
        splitter = self._get_splitter(chunk_size, chunk_overlap)

        # Split into chunks
        try:
            chunks = splitter.chunks(full_text)
        except Exception as exc:
            logger.error("Chunking failed strategy=%s: %s", strategy, exc)
            raise

        # Build chunk -> page mapping
        chunk_page_map = self._map_chunks_to_pages(chunks, full_text, page_offsets)

        # Record metrics
        CHUNKS_CREATED_TOTAL.labels(
            library_type=library_type,
            strategy=strategy,
        ).inc(len(chunks))

        for chunk_text in chunks:
            CHUNK_SIZE_TOKENS.observe(len(chunk_text.split()))

        logger.info(
            "Chunked document strategy=%s chunks=%d avg_size=%d",
            strategy,
            len(chunks),
            sum(len(c) for c in chunks) // max(len(chunks), 1),
        )

        return ChunkResult(
            chunks=chunks,
            chunk_page_map=chunk_page_map,
            strategy=strategy,
        )

    def _get_splitter(self, chunk_size: int, chunk_overlap: int):
        """
        Get or create a semantic text splitter for the given parameters.

        :param chunk_size: Maximum chunk size in tokens.
        :param chunk_overlap: Overlap between chunks in tokens.
        :returns: TextSplitter instance.
        """
        cache_key = (chunk_size, chunk_overlap)
        if cache_key in self._splitter_cache:
            return self._splitter_cache[cache_key]

        from semantic_text_splitter import TextSplitter
        from tokenizers import Tokenizer

        try:
            tokenizer = Tokenizer.from_pretrained(self._tokenizer_name)
            splitter = TextSplitter.from_huggingface_tokenizer(
                tokenizer,
                capacity=chunk_size,
                overlap=chunk_overlap,
            )
            logger.debug(
                "Created text splitter tokenizer=%s capacity=%d overlap=%d",
                self._tokenizer_name,
                chunk_size,
                chunk_overlap,
            )
        except Exception as exc:
            logger.warning(
                "Failed to load tokenizer %s: %s, falling back to %s",
                self._tokenizer_name,
                exc,
                DEFAULT_TOKENIZER,
            )
            tokenizer = Tokenizer.from_pretrained(DEFAULT_TOKENIZER)
            splitter = TextSplitter.from_huggingface_tokenizer(
                tokenizer,
                capacity=chunk_size,
                overlap=chunk_overlap,
            )

        self._splitter_cache[cache_key] = splitter
        return splitter

    def _combine_text_blocks(
        self, text_blocks: list[TextBlock]
    ) -> tuple[str, list[tuple[int, int]]]:
        """
        Combine text blocks into a single string, tracking page offsets.

        :param text_blocks: List of TextBlock from parser.
        :returns: Tuple of (combined_text, page_offsets) where page_offsets
                  is a list of (char_offset, page_number).
        """
        parts: list[str] = []
        page_offsets: list[tuple[int, int]] = []
        current_offset = 0

        for block in text_blocks:
            page_offsets.append((current_offset, block.page))
            parts.append(block.text)
            current_offset += len(block.text) + 2  # +2 for paragraph separator

        return "\n\n".join(parts), page_offsets

    def _map_chunks_to_pages(
        self,
        chunks: list[str],
        full_text: str,
        page_offsets: list[tuple[int, int]],
    ) -> dict[int, int]:
        """
        Map each chunk index to its source page number.

        :param chunks: List of chunk strings.
        :param full_text: Combined document text.
        :param page_offsets: List of (char_offset, page_number).
        :returns: Dict mapping chunk_index -> page_number.
        """
        chunk_page_map: dict[int, int] = {}
        search_start = 0

        for chunk_idx, chunk_text in enumerate(chunks):
            # Find where this chunk starts in the full text
            pos = full_text.find(chunk_text[:100], search_start)
            if pos == -1:
                pos = search_start

            # Find which page this position belongs to
            page = 0
            for offset, page_num in page_offsets:
                if offset <= pos:
                    page = page_num
                else:
                    break

            chunk_page_map[chunk_idx] = page
            search_start = max(search_start, pos)

        return chunk_page_map
