from __future__ import annotations

import re
from typing import Any, Dict, List, Sequence
from uuid import uuid4

from idxr.vectorize_lib.utils import count_tokens, get_token_encoder
from idxr.vectorize_lib.compact.document_compactor import CHROMA_DOCUMENT_SIZE_LIMIT

from ..interfaces import ChunkingStrategy
from ..models import Chunk, DocumentIngestContext, DocumentIngestRequest

# Pattern to match page markers inserted by PDF extraction
PAGE_MARKER_PATTERN = re.compile(r"<<PAGE:(\d+)>>\n?")


class TokenCountingChunker(ChunkingStrategy):
    """
    Splits text into token-aware windows with optional overlap.
    """

    def __init__(
        self,
        *,
        encoder_model: str,
        default_chunk_size: int = 800,
        default_overlap: int = 200,
        byte_limit: int = CHROMA_DOCUMENT_SIZE_LIMIT,
    ) -> None:
        self._encoder_model = encoder_model
        self._default_chunk_size = default_chunk_size
        self._default_overlap = default_overlap
        self._byte_limit = byte_limit
        self._encoder = get_token_encoder(encoder_model)

    def chunk(
        self,
        text: str,
        request: DocumentIngestRequest,
        context: DocumentIngestContext,
    ) -> Sequence[Chunk]:
        chunk_size = max(1, request.chunk_size_tokens or self._default_chunk_size)
        overlap = max(
            0,
            min(
                chunk_size // 2,
                request.chunk_overlap_tokens or self._default_overlap,
            ),
        )

        # Extract page mapping before tokenization (for PDFs)
        page_map = self._build_page_map(text)

        # Remove page markers from text before chunking
        clean_text = PAGE_MARKER_PATTERN.sub("", text)

        tokens = self._encoder.encode(clean_text)
        total_tokens = len(tokens)
        if total_tokens == 0:
            return []

        chunks: List[Chunk] = []
        start = 0
        chunk_no = 0
        while start < total_tokens:
            end = min(total_tokens, start + chunk_size)
            chunk_tokens = tokens[start:end]
            chunk_text = self._encoder.decode(chunk_tokens)
            chunk_bytes = len(chunk_text.encode("utf-8"))
            if chunk_bytes > self._byte_limit:
                chunk_text = self._hard_trim(chunk_text, self._byte_limit)
                chunk_tokens = self._encoder.encode(chunk_text)
                chunk_bytes = len(chunk_text.encode("utf-8"))

            # Determine page number(s) for this chunk
            page_no = self._find_page_for_chunk(chunk_text, page_map, clean_text)

            chunk_id = f"{context.conversation_id}:{chunk_no}:{uuid4().hex[:8]}"
            chunk_metadata: Dict[str, Any] = {
                "chunk_no": chunk_no,  # Renamed from chunk_index for consistency
                "chunk_offset_tokens": start,
                "token_count": count_tokens(chunk_text, self._encoder),
                "byte_count": chunk_bytes,
            }

            # Add page_no for PDFs
            if page_no is not None:
                chunk_metadata["page_no"] = page_no

            # Add source_url for all documents (important for website chunks)
            if request.source_url:
                chunk_metadata["url"] = request.source_url

            chunks.append(
                Chunk(
                    chunk_id=chunk_id,
                    text=chunk_text,
                    token_count=chunk_metadata["token_count"],
                    byte_count=chunk_bytes,
                    metadata=chunk_metadata,
                )
            )

            if end >= total_tokens:
                break
            start = max(0, end - overlap)
            chunk_no += 1
        return chunks

    def _build_page_map(self, text: str) -> List[tuple[int, int]]:
        """
        Build a map of (char_offset, page_number) from page markers.

        Returns a list of tuples where each tuple is (character offset, page number).
        This allows us to determine which page a chunk belongs to.
        """
        page_map: List[tuple[int, int]] = []
        for match in PAGE_MARKER_PATTERN.finditer(text):
            page_num = int(match.group(1))
            page_map.append((match.start(), page_num))
        return page_map

    def _find_page_for_chunk(
        self, chunk_text: str, page_map: List[tuple[int, int]], clean_text: str
    ) -> int | None:
        """
        Determine which page a chunk belongs to based on page markers.

        If the chunk spans multiple pages, returns the starting page number.
        Returns None if no page markers were found (non-PDF documents).
        """
        if not page_map:
            return None

        # Find where this chunk appears in the clean text
        chunk_start = clean_text.find(
            chunk_text[:100] if len(chunk_text) > 100 else chunk_text
        )
        if chunk_start == -1:
            # Fallback: use the first page if we can't find the chunk
            return page_map[0][1] if page_map else None

        # Find the last page marker before this chunk
        current_page = page_map[0][1]  # Default to first page
        for offset, page_num in page_map:
            if offset <= chunk_start:
                current_page = page_num
            else:
                break

        return current_page

    def _hard_trim(self, text: str, byte_limit: int) -> str:
        encoded = text.encode("utf-8")
        if len(encoded) <= byte_limit:
            return text
        return encoded[:byte_limit].decode("utf-8", errors="ignore").rstrip()
