"""Text chunking strategies for RAG pipelines."""

from __future__ import annotations

import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field


@dataclass
class Chunk:
    """A chunk of text with metadata."""

    text: str
    index: int
    metadata: dict = field(default_factory=dict)

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

    def __str__(self) -> str:
        return self.text


class ChunkingStrategy(ABC):
    """Base class for chunking strategies."""

    @abstractmethod
    def chunk(self, text: str, metadata: dict | None = None) -> list[Chunk]:
        ...


class FixedSizeChunker(ChunkingStrategy):
    """Split text into fixed-size character chunks with optional overlap."""

    def __init__(self, chunk_size: int = 500, overlap: int = 50):
        if overlap >= chunk_size:
            raise ValueError(f"Overlap ({overlap}) must be less than chunk_size ({chunk_size})")
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk(self, text: str, metadata: dict | None = None) -> list[Chunk]:
        if not text:
            return []

        meta = metadata or {}
        chunks = []
        start = 0
        index = 0

        while start < len(text):
            end = start + self.chunk_size
            chunk_text = text[start:end]
            chunks.append(Chunk(text=chunk_text, index=index, metadata={**meta}))
            index += 1
            start += self.chunk_size - self.overlap

        return chunks


class SentenceChunker(ChunkingStrategy):
    """Split text into chunks of N sentences."""

    def __init__(self, max_sentences: int = 3):
        self.max_sentences = max_sentences

    def chunk(self, text: str, metadata: dict | None = None) -> list[Chunk]:
        if not text.strip():
            return []

        meta = metadata or {}
        sentences = re.split(r'(?<=[.!?])\s+', text.strip())
        sentences = [s for s in sentences if s.strip()]

        chunks = []
        for i in range(0, len(sentences), self.max_sentences):
            group = sentences[i:i + self.max_sentences]
            chunk_text = " ".join(group)
            chunks.append(Chunk(text=chunk_text, index=len(chunks), metadata={**meta}))

        return chunks


class RecursiveChunker(ChunkingStrategy):
    """Recursively split text using a hierarchy of separators.

    Tries paragraph breaks first, then sentences, then fixed-size.
    """

    SEPARATORS = ["\n\n", "\n", ". ", " "]

    def __init__(self, chunk_size: int = 500, overlap: int = 50):
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk(self, text: str, metadata: dict | None = None) -> list[Chunk]:
        if not text.strip():
            return []

        meta = metadata or {}
        raw_chunks = self._split(text, self.SEPARATORS)

        chunks = []
        for i, raw in enumerate(raw_chunks):
            chunks.append(Chunk(text=raw, index=i, metadata={**meta}))
        return chunks

    def _split(self, text: str, separators: list[str]) -> list[str]:
        if len(text) <= self.chunk_size:
            return [text.strip()] if text.strip() else []

        if not separators:
            # Fallback to fixed-size split
            result = []
            start = 0
            while start < len(text):
                result.append(text[start:start + self.chunk_size].strip())
                start += self.chunk_size - self.overlap
            return [r for r in result if r]

        sep = separators[0]
        parts = text.split(sep)

        # Try to merge parts into chunks that fit
        merged = []
        current = ""

        for part in parts:
            candidate = (current + sep + part).strip() if current else part.strip()
            if len(candidate) <= self.chunk_size:
                current = candidate
            else:
                if current:
                    merged.append(current)
                # If part itself exceeds chunk_size, recurse with finer separator
                if len(part.strip()) > self.chunk_size:
                    merged.extend(self._split(part, separators[1:]))
                elif part.strip():
                    current = part.strip()
                else:
                    current = ""

        if current:
            merged.append(current)

        return [m for m in merged if m]
