import logging
from typing import List, Optional

from fine_chunker.engine.base import BaseEngine
from fine_chunker.model.base import BaseChunkingModel
from fine_chunker.preprocessing import DefaultPreprocessor, Preprocessor
from fine_chunker.types import Chunk

logger = logging.getLogger("fine_chunker.engine.recursive")


class RecursiveEngine(BaseEngine):
    """
    Semantic chunking engine that recursively splits text using a ModernBERT model.

    The RecursiveEngine is designed to handle large text inputs by breaking them into
    smaller, semantically meaningful chunks. It uses a recursive approach to ensure
    that chunks are within a specified size range, leveraging a ModernBERT-based model
    to predict optimal split points.

    Attributes:
        model (BaseChunkingModel): The ModernBERT-based model used for chunking.
        preprocessor (Optional[Preprocessor]): Preprocessor for cleaning and preparing text.
        threshold_start (float): Initial threshold for boundary prediction.
        threshold_step (float): Step by which the threshold decreases at each recursion level.
        max_depth (int): Maximum recursion depth for splitting.
        min_chunk_size (int): Minimum size of a chunk in characters.
        max_chunk_size (int): Maximum size of a chunk in characters. (Can be exceed in some cases.)

    Methods:
        chunk(text: str) -> List[Chunk]:
            Splits the input text into a list of chunks.
    """

    def __init__(
        self,
        model: BaseChunkingModel,
        preprocessor: Optional[Preprocessor] = None,
        threshold_start: float = 0.5,
        threshold_step: float = 0.1,
        max_depth: int = 3,
        min_chunk_size: int = 350,
        max_chunk_size: int = 1000,
    ):
        self.model = model
        self.preprocessor = preprocessor or DefaultPreprocessor()
        self.threshold_start = threshold_start
        self.threshold_step = threshold_step
        self.max_depth = max_depth
        self.min_chunk_size = min_chunk_size
        self.max_chunk_size = max_chunk_size
        self._validate_args()

        logger.debug(
            "RecursiveEngine initialized | depth=%d threshold=%.2f -> %.2f",
            max_depth,
            threshold_start,
            threshold_start - (max_depth * threshold_step),
        )

    def chunk(self, text: str) -> List[Chunk]:
        clean = self.preprocessor.preprocess(text)
        if not clean:
            return []

        windows = self._get_text_windows(clean)

        all_fragments: List[str] = []
        for i, window in enumerate(windows):
            if len(windows) > 1:
                logger.info(
                    "Processing window %d/%d (%d chars)",
                    i + 1,
                    len(windows),
                    len(window),
                )

            fragments = self._split(window, self.threshold_start, depth=0)
            all_fragments.extend(fragments)

        chunks = [Chunk(index=i, content=f) for i, f in enumerate(all_fragments)]
        logger.info("Chunking complete | total_chunks=%d", len(chunks))
        return chunks

    def _get_text_windows(self, text: str) -> List[str]:
        """
        Splits text into strings that are safe for the model to process (within max_length).
        Uses the model to find a semantic split point in the 'safe zone' at the end of each window.
        """
        all_ids = self.model.tokenize(text, add_special_tokens=False)
        max_len = self.model.max_length() - 2  # buffer for special tokens

        if len(all_ids) <= max_len:
            return [text]

        logger.info(
            "Input exceeds model limit (%d > %d tokens). Splitting into windows... - it may take some time.",
            len(all_ids),
            max_len,
        )

        windows: List[str] = []
        curr_pos = 0
        safe_zone_size = 500

        while curr_pos < len(all_ids):
            end_pos = min(curr_pos + max_len, len(all_ids))

            if end_pos == len(all_ids):
                windows.append(self.model.decode(all_ids[curr_pos:]))
                break

            window_ids = all_ids[curr_pos:end_pos]
            predictions = self.model.predict(window_ids, self.threshold_start)[
                "predictions"
            ]

            split_idx = -1
            zone_start = max(0, len(window_ids) - safe_zone_size)

            for i in range(len(window_ids) - 1, zone_start, -1):
                if predictions[i].item() == 1:
                    split_idx = i
                    break

            if split_idx == -1:
                logger.debug(
                    "No boundary in safe zone at token %d, forcing split.",
                    curr_pos + end_pos,
                )
                split_idx = len(window_ids) - 1

            windows.append(self.model.decode(window_ids[:split_idx]))
            curr_pos += split_idx

        return windows

    def _split(self, text: str, threshold: float, depth: int) -> List[str]:
        """Recursive split logic."""
        raw_fragments = self._model_split(text, threshold)

        result: List[str] = []
        for fragment in raw_fragments:
            flen = len(fragment)

            if flen > self.max_chunk_size and depth < self.max_depth:
                next_threshold = threshold - self.threshold_step
                if next_threshold > 0:
                    sub = self._split(fragment, next_threshold, depth + 1)
                    result.extend(sub)
                    continue

            result.append(fragment)

        return self._merge_small(result)

    def _model_split(self, text: str, threshold: float) -> List[str]:
        """Runs the model and returns fragments based on predicted boundaries."""
        result = self.model.tokenize_and_predict(text, threshold)
        input_ids = result["input_ids"]
        predictions = result["predictions"]

        fragments: List[str] = []
        current_token_ids: List[int] = []

        # Start from 1 to skip [CLS]
        for i in range(1, len(input_ids)):
            token_id = input_ids[i].item()
            is_boundary = predictions[i].item() == 1

            if is_boundary and current_token_ids:
                frag = self.model.decode(current_token_ids)
                if frag:
                    fragments.append(frag)
                current_token_ids = [token_id]
            else:
                current_token_ids.append(token_id)

        if current_token_ids:
            frag = self.model.decode(current_token_ids)
            if frag:
                fragments.append(frag)

        return fragments if fragments else [text]

    def _merge_small(self, fragments: List[str]) -> List[str]:
        """Merges fragments smaller than min_chunk_size."""
        if not fragments:
            return fragments

        merged: List[str] = [fragments[0]]

        for fragment in fragments[1:]:
            would_be_size = len(merged[-1]) + 1 + len(fragment)
            if would_be_size <= self.max_chunk_size:
                merged[-1] = merged[-1] + " " + fragment
            else:
                merged.append(fragment)

        return merged

    def _validate_args(self):
        if self.threshold_start <= 0 or self.threshold_start >= 1:
            raise ValueError("threshold_start must be in the range (0, 1)")
        if self.threshold_step <= 0 or self.threshold_step >= 1:
            raise ValueError("threshold_step must be in the range (0, 1)")
        if self.max_depth < 0:
            raise ValueError("max_depth must be non-negative")
        if self.min_chunk_size <= 0:
            raise ValueError("min_chunk_size must be positive")
        if self.max_chunk_size <= 0:
            raise ValueError("max_chunk_size must be positive")
