"""Module containing the SlumberChunker."""

import re
from bisect import bisect_left
from itertools import accumulate
from typing import Literal, Optional, Union

import chonkie_core
from tqdm import tqdm

from chonkie.genie import BaseGenie, GeminiGenie
from chonkie.logger import get_logger
from chonkie.pipeline import chunker
from chonkie.tokenizer import TokenizerProtocol
from chonkie.types import Chunk, RecursiveLevel, RecursiveRules

from .base import BaseChunker

logger = get_logger(__name__)


JSON_PROMPT_TEMPLATE = """<task> You are given a set of texts between the starting tag <passages> and ending tag </passages>. Each text is labeled as 'ID `N`' where 'N' is the passage number. Your task is to find the first passage where the content clearly separates from the previous passages in topic and/or semantics. </task>

<rules>
Follow the following rules while finding the splitting passage:
- Always return the answer as a JSON parsable object with the 'split_index' key having a value of the first passage where the topic changes.
- Avoid very long groups of paragraphs. Aim for a good balance between identifying content shifts and keeping groups manageable.
- If no clear `split_index` is found, return N + 1, where N is the index of the last passage.
</rules>

<passages>
{passages}
</passages>
"""

TEXT_PROMPT_TEMPLATE = """<context>
You are chunking text for a RAG (Retrieval Augmented Generation) system.
Good chunks should be topically coherent so that relevant information can be retrieved together.
</context>

<task>
Find the first passage ID where the topic shifts enough to start a new chunk.
</task>

<examples>
Example: If passages discuss "philosophy" then suddenly "cooking recipes", split there.
Example: If passages all discuss aspects of the same topic, keep them together.
</examples>

<passages>
{passages}
</passages>

<format>
Return ONLY the ID number, or N+1 (where N is the last passage ID) if all passages should stay together.
Do not include any explanation or additional text.
</format>

<answer>"""


@chunker("slumber")
class SlumberChunker(BaseChunker):
    """SlumberChunker is a chunker based on the LumberChunker — but slightly different."""

    def __init__(
        self,
        genie: Optional[BaseGenie] = None,
        tokenizer: Union[str, TokenizerProtocol] = "character",
        chunk_size: int = 2048,
        rules: Optional[RecursiveRules] = None,
        candidate_size: int = 128,
        min_characters_per_chunk: int = 24,
        extract_mode: Literal["text", "json", "auto"] = "auto",
        max_retries: int = 3,
        verbose: bool = True,
    ):
        """Initialize the SlumberChunker.

        Args:
            genie (Optional[BaseGenie]): The genie to use.
            tokenizer: The tokenizer to use.
            chunk_size (int): The size of the chunks to create.
            rules (Optional[RecursiveRules]): The rules to use to split the candidate chunks.
            candidate_size (int): The size of the candidate splits that the chunker will consider.
            min_characters_per_chunk (int): The minimum number of characters per chunk.
            extract_mode: Mode for extracting split index from LLM response.
                - "json": Use structured JSON output via generate_json() (requires genie support)
                - "text": Use plain text generation via generate() and parse integer response
                - "auto": Auto-detect based on genie capabilities (default)
            max_retries (int): Maximum retries for text mode parsing failures.
            verbose (bool): Whether to print verbose output.

        """
        # Since the BaseChunker sets and defines the tokenizer for us, we don't have to worry.
        super().__init__(tokenizer)

        # If the genie is not provided, use the default GeminiGenie
        if genie is None:
            genie = GeminiGenie()

        self.genie = genie
        self.max_retries = max_retries

        # Determine effective extract_mode
        self.extract_mode = self._determine_extract_mode(extract_mode)

        # Lazy import pydantic only when using JSON mode
        self.Split: Optional[type] = None
        if self.extract_mode == "json":
            try:
                from pydantic import BaseModel
            except ImportError as ie:
                raise ImportError(
                    "The SlumberChunker requires the pydantic library for extract_mode='json'. "
                    "Please install it using `pip install chonkie[genie]` or use extract_mode='text'.",
                ) from ie

            class Split(BaseModel):
                split_index: int

            self.Split = Split

        # Set the parameters for the SlumberChunker
        self.chunk_size = chunk_size
        self.candidate_size = candidate_size
        self.rules = rules if rules is not None else RecursiveRules()
        self.min_characters_per_chunk = min_characters_per_chunk
        self.verbose = verbose

        # Set the template based on extract_mode
        if self.extract_mode == "json":
            self.template = JSON_PROMPT_TEMPLATE
        else:
            self.template = TEXT_PROMPT_TEMPLATE
        self.sep = "✄"
        self._CHARS_PER_TOKEN = 6.5

        # Set the _use_multiprocessing to False, since we don't know the
        # behaviour of the Genie under multiprocessing conditions
        self._use_multiprocessing = False

    def _determine_extract_mode(
        self, mode: Optional[Literal["text", "json", "auto"]]
    ) -> Literal["text", "json"]:
        """Determine the effective extract mode based on genie capabilities.

        Args:
            mode: The requested extract mode ("text", "json", or "auto")

        Returns:
            The effective extract mode to use ("text" or "json")

        """
        if mode == "json":
            return "json"
        elif mode == "text":
            return "text"
        elif mode == "auto" or mode is None:
            # Auto-detect based on genie capabilities
            return self._detect_genie_json_support()
        else:
            raise ValueError(f"Invalid extract_mode: {mode}. Must be 'text', 'json', or 'auto'.")

    def _detect_genie_json_support(self) -> Literal["text", "json"]:
        """Detect if the genie supports JSON generation.

        Returns:
            "json" if genie supports generate_json, "text" otherwise

        """
        try:
            # Get the generate_json method from the genie's class
            genie_method = type(self.genie).generate_json
            base_method = BaseGenie.generate_json

            # If the method is the same as the base class (not overridden), use text
            if genie_method is base_method:
                logger.debug(
                    f"Genie {type(self.genie).__name__} does not override generate_json, using text mode"
                )
                return "text"

            # Method is overridden, assume JSON support
            logger.debug(
                f"Genie {type(self.genie).__name__} overrides generate_json, using JSON mode"
            )
            return "json"
        except Exception as e:
            logger.warning(f"Failed to detect genie JSON support: {e}, defaulting to text mode")
            return "text"

    def _extract_index_from_text(self, response: str) -> int:
        """Extract the integer index from a plain text response.

        Args:
            response: The raw text response from the genie

        Returns:
            The extracted split index as an integer

        Raises:
            ValueError: If no valid integer can be extracted

        """
        # Clean the response
        cleaned = response.strip()

        # Try direct integer parse first
        try:
            return int(cleaned)
        except ValueError:
            pass

        # Try to find any integer in the response
        match = re.search(r"(\d+)", cleaned)
        if match:
            return int(match.group(1))

        raise ValueError(f"Could not extract integer from response: '{response}'")

    def _get_split_index(self, prompt: str, current_pos: int, group_end_index: int) -> int:
        """Get the split index from the genie using the appropriate extraction mode.

        Args:
            prompt: The formatted prompt to send to the genie
            current_pos: The current position in the splits list
            group_end_index: The end index of the current group (fallback if extraction fails)

        Returns:
            The predicted split index

        """
        if self.extract_mode == "json":
            return self._get_split_index_json(prompt, group_end_index)
        else:
            return self._get_split_index_text(prompt, group_end_index)

    def _get_split_index_json(self, prompt: str, group_end_index: int) -> int:
        """Get split index using JSON extraction mode with retry logic.

        Args:
            prompt: The formatted prompt
            group_end_index: End index of current group (used for fallback)

        Returns:
            The split index from the JSON response

        """
        last_error = None

        for attempt in range(self.max_retries):
            try:
                response = self.genie.generate_json(prompt, self.Split)
                index = int(response["split_index"])
                if index > group_end_index:
                    raise ValueError(
                        f"Split index {index} is out of bounds (max {group_end_index})"
                    )
                return index
            except (KeyboardInterrupt, SystemExit):
                raise
            except Exception as e:
                last_error = e
                logger.debug(
                    f"JSON extraction attempt {attempt + 1}/{self.max_retries} failed: {e}"
                )
                continue

        # All retries failed - keep passages together by returning group end
        logger.debug(
            f"JSON extraction failed after {self.max_retries} attempts. "
            f"Last error: {last_error}. Keeping passages together. Using fallback index {group_end_index}."
        )
        return group_end_index

    def _get_split_index_text(self, prompt: str, group_end_index: int) -> int:
        """Get split index using text extraction mode with retry logic.

        Args:
            prompt: The formatted prompt
            group_end_index: End index of current group (used for fallback)

        Returns:
            The extracted split index

        """
        last_error = None

        for attempt in range(self.max_retries):
            try:
                response = self.genie.generate(prompt)
                index = self._extract_index_from_text(response)
                if index > group_end_index:
                    raise ValueError(
                        f"Split index {index} is out of bounds (max {group_end_index})"
                    )
                return index
            except (KeyboardInterrupt, SystemExit):
                raise
            except Exception as e:
                last_error = e
                logger.debug(
                    f"Text extraction attempt {attempt + 1}/{self.max_retries} failed: {e}"
                )
                continue

        # All retries failed - keep passages together by returning group end
        logger.debug(
            f"Text extraction failed after {self.max_retries} attempts. "
            f"Last error: {last_error}. Keeping passages together. Using fallback index {group_end_index}."
        )
        return group_end_index

    def _split_text(self, text: str, recursive_level: RecursiveLevel) -> list[str]:
        """Split the text into chunks using the delimiters."""
        if recursive_level.delimiters:
            include_mode = recursive_level.include_delim or "prev"
            text_bytes = text.encode("utf-8")

            # Check if we have multi-byte delimiters
            delimiters = recursive_level.delimiters
            if isinstance(delimiters, str):
                delimiters = [delimiters]

            has_multibyte = any(len(d) > 1 for d in delimiters)

            if has_multibyte:
                # Use split_pattern_offsets for multi-byte patterns
                patterns = [d.encode("utf-8") for d in delimiters]
                offsets = chonkie_core.split_pattern_offsets(
                    text_bytes,
                    patterns=patterns,
                    include_delim=include_mode,
                    min_chars=self.min_characters_per_chunk,
                )
            else:
                # Use faster split_offsets for single-byte delimiters
                delim_bytes = "".join(delimiters).encode("utf-8")
                offsets = chonkie_core.split_offsets(
                    text_bytes,
                    delimiters=delim_bytes,
                    include_delim=include_mode,
                    min_chars=self.min_characters_per_chunk,
                )

            splits = [text_bytes[start:end].decode("utf-8") for start, end in offsets]
            return [s for s in splits if s]
        elif recursive_level.whitespace:
            # Split on whitespace using split_offsets (preserves spaces for reconstruction)
            text_bytes = text.encode("utf-8")
            include_mode = recursive_level.include_delim or "prev"
            offsets = chonkie_core.split_offsets(
                text_bytes,
                delimiters=b" ",
                include_delim=include_mode,
                min_chars=self.min_characters_per_chunk,
            )
            splits = [text_bytes[start:end].decode("utf-8") for start, end in offsets]
            return [s for s in splits if s]
        else:
            # Token-based splitting
            encoded = self.tokenizer.encode(text)
            token_splits = [
                encoded[i : i + self.chunk_size] for i in range(0, len(encoded), self.chunk_size)
            ]
            return list(self.tokenizer.decode_batch(token_splits))

    def _recursive_split(self, text: str, level: int = 0, offset: int = 0) -> list[Chunk]:
        """Recursively split the text into chunks."""
        if not self.rules.levels or level >= len(self.rules.levels):
            return [
                Chunk(
                    text=text,
                    start_index=offset,
                    end_index=offset + len(text),
                    token_count=self.tokenizer.count_tokens(text),
                ),
            ]

        # Do the first split based on the level provided
        splits = self._split_text(text, self.rules.levels[level])

        # Calculate the token_count of each of the splits
        token_counts = self.tokenizer.count_tokens_batch(splits)

        # Loop throught the splits to see if any split
        chunks = []
        current_offset = offset
        for split, token_count in zip(splits, token_counts):
            # If the token_count is more than the self.candidate_size,
            # then call the recursive_split function on it with a higher level
            if token_count > self.candidate_size:
                child_chunks = self._recursive_split(split, level + 1, current_offset)
                chunks.extend(child_chunks)
            else:
                chunks.append(
                    Chunk(
                        text=split,
                        start_index=current_offset,
                        end_index=current_offset + len(split),
                        token_count=token_count,
                    ),
                )

            # Add the offset as the length of the split
            current_offset += len(split)

        return chunks

    def _prepare_splits(self, splits: list[Chunk]) -> list[str]:
        """Prepare the splits for the chunker."""
        return [
            f"ID {i}: " + split.text.replace("\n", " ").strip() for (i, split) in enumerate(splits)
        ]

    def _get_cumulative_token_counts(self, splits: list[Chunk]) -> list[int]:
        """Get the cumulative token counts for the splits."""
        return list(accumulate([0] + [split.token_count for split in splits]))

    def chunk(self, text: str) -> list[Chunk]:
        """Chunk the text with the SlumberChunker."""
        logger.debug(f"Starting slumber chunking for text of length {len(text)}")

        # Store original text for accurate extraction
        original_text = text

        splits = self._recursive_split(text, level=0, offset=0)
        logger.debug(
            f"Created {len(splits)} initial splits for LLM-based semantic boundary detection",
        )

        # Add the IDS to the splits
        prepared_split_texts = self._prepare_splits(splits)

        # Calculate the cumulative token counts for each split
        cumulative_token_counts = self._get_cumulative_token_counts(splits)

        # If self.verbose has been set to True, show a TQDM progress bar for the text
        if self.verbose:
            progress_bar = tqdm(
                total=len(splits),
                desc="🦛",
                unit="split",
                bar_format="{desc} ch{bar:20}nk {percentage:3.0f}% • {n_fmt}/{total_fmt} splits processed [{elapsed}<{remaining}, {rate_fmt}] 🌱",
                ascii=" o",
            )

        # Pass the self.chunk_size amount of context through the Genie,
        # so we can control how much context the Genie gets as well.
        # This is especially useful for models that don't have long context
        # or exhibit weakend reasoning ability over longer texts.
        chunks = []
        current_pos = 0
        current_token_count = 0
        while current_pos < len(splits):
            # bisect_left can return 0? No because input_size > 0 and first value is 0
            group_end_index = min(
                bisect_left(cumulative_token_counts, current_token_count + self.chunk_size) - 1,
                len(splits),
            )

            if group_end_index == current_pos:
                group_end_index += 1

            prompt = self.template.format(
                passages="\n".join(prepared_split_texts[current_pos:group_end_index]),
            )
            # response is always <= group_end_index
            response = self._get_split_index(prompt, current_pos, group_end_index)

            # Make sure that the response doesn't bug out and return a index smaller
            # than the current position
            if current_pos >= response:
                response = current_pos + 1

            # Extract text directly from original source to preserve all spacing and formatting
            start_idx = splits[current_pos].start_index
            end_idx = splits[response - 1].end_index

            chunks.append(
                Chunk(
                    text=original_text[start_idx:end_idx],
                    start_index=start_idx,
                    end_index=end_idx,
                    token_count=sum([split.token_count for split in splits[current_pos:response]]),
                ),
            )

            current_token_count = cumulative_token_counts[response]
            current_pos = response

            if self.verbose:
                progress_bar.update(current_pos - progress_bar.n)

        logger.info(f"Created {len(chunks)} chunks using LLM-guided semantic splitting")
        return chunks

    def __repr__(self) -> str:
        """Return a string representation of the SlumberChunker."""
        return (
            f"SlumberChunker(genie={self.genie}, "
            f"tokenizer={self.tokenizer}, "
            f"chunk_size={self.chunk_size}, "
            f"candidate_size={self.candidate_size}, "
            f"min_characters_per_chunk={self.min_characters_per_chunk}, "
            f"extract_mode={self.extract_mode}, "
            f"max_retries={self.max_retries})"
        )
