"""Text splitting strategies for RAG system."""
from __future__ import annotations

from abc import ABC, abstractmethod

from logos_adk.rag import Chunk, Document


class TextSplitter(ABC):
    """Abstract base class for text splitters."""

    @abstractmethod
    def split(self, document: Document) -> list[Chunk]:
        """Split a document into chunks."""
        ...


class RecursiveCharacterSplitter(TextSplitter):
    """Recursively splits text by separators with overlap."""

    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        separators: list[str] | None = None,
    ) -> None:
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = separators or ["\n\n", "\n", ". ", " "]

    def split(self, document: Document) -> list[Chunk]:
        text = document.content.strip()
        if not text:
            return []

        pieces = self._split_text(text, self.separators)
        chunks: list[Chunk] = []

        for i, piece in enumerate(pieces):
            chunks.append(Chunk(
                content=piece,
                metadata={**document.metadata, "chunk_index": i},
                document_id=document.id,
            ))

        return chunks

    def _split_text(self, text: str, separators: list[str]) -> list[str]:
        """Recursively split text using separators."""
        if len(text) <= self.chunk_size:
            return [text]

        for sep in separators:
            if sep in text:
                parts = text.split(sep)
                return self._merge_parts(parts, sep, separators)

        return self._force_split(text)

    def _merge_parts(
        self, parts: list[str], sep: str, separators: list[str]
    ) -> list[str]:
        """Merge small parts together, split large parts recursively."""
        result: list[str] = []
        current: list[str] = []
        current_len = 0

        for part in parts:
            part = part.strip()
            if not part:
                continue

            part_len = len(part)
            sep_len = len(sep) if current else 0

            if current_len + sep_len + part_len <= self.chunk_size:
                current.append(part)
                current_len += sep_len + part_len
            else:
                if current:
                    result.append(sep.join(current).strip())
                    if self.chunk_overlap > 0:
                        overlap_text = sep.join(current).strip()
                        overlap_start = max(0, len(overlap_text) - self.chunk_overlap)
                        overlap = overlap_text[overlap_start:].strip()
                        current = [overlap, part] if overlap else [part]
                        current_len = len(sep.join(current))
                    else:
                        current = [part]
                        current_len = part_len
                else:
                    remaining_seps = separators[separators.index(sep) + 1:] if sep in separators else []
                    if remaining_seps:
                        result.extend(self._split_text(part, remaining_seps))
                    else:
                        result.extend(self._force_split(part))
                    current = []
                    current_len = 0

        if current:
            result.append(sep.join(current).strip())

        return [r for r in result if r]

    def _force_split(self, text: str) -> list[str]:
        """Force split by chunk_size when no separator works."""
        result: list[str] = []
        start = 0
        while start < len(text):
            end = start + self.chunk_size
            result.append(text[start:end].strip())
            start = end - self.chunk_overlap if self.chunk_overlap > 0 else end
        return [r for r in result if r]
