"""
cogniq.chunkers — Smart document chunking.

SmartChunker detects semantic boundaries (headings, paragraph breaks,
sentence endings with topic shifts) rather than blindly splitting by tokens.
"""

from __future__ import annotations
import re
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional


class BaseChunker(ABC):
    @abstractmethod
    def chunk(self, text: str, source: Optional[str] = None) -> List[Dict[str, Any]]: ...

    def __repr__(self):
        return f"{self.__class__.__name__}()"


# ─────────────────────────────────────────────
# FIXED SIZE
# ─────────────────────────────────────────────

class FixedChunker(BaseChunker):
    """Split by word count with overlap."""

    def __init__(self, size: int = 300, overlap: int = 50):
        self.size = size
        self.overlap = overlap

    def chunk(self, text: str, source: Optional[str] = None) -> List[Dict[str, Any]]:
        words = text.split()
        chunks = []
        step = max(1, self.size - self.overlap)
        for i in range(0, max(1, len(words) - self.overlap), step):
            chunk_words = words[i: i + self.size]
            if not chunk_words:
                break
            chunks.append({
                "text": " ".join(chunk_words),
                "source": source,
                "chunk_idx": len(chunks),
                "start_word": i,
            })
        return chunks


# ─────────────────────────────────────────────
# SENTENCE
# ─────────────────────────────────────────────

class SentenceChunker(BaseChunker):
    """Group sentences into chunks of ~target_sentences each."""

    def __init__(self, target_sentences: int = 5, overlap_sentences: int = 1):
        self.target = target_sentences
        self.overlap = overlap_sentences

    def _split_sentences(self, text: str) -> List[str]:
        # Simple but effective sentence splitter
        parts = re.split(r'(?<=[.!?])\s+(?=[A-Z"\'])', text)
        return [p.strip() for p in parts if p.strip()]

    def chunk(self, text: str, source: Optional[str] = None) -> List[Dict[str, Any]]:
        sentences = self._split_sentences(text)
        if not sentences:
            return [{"text": text, "source": source, "chunk_idx": 0}]

        chunks = []
        step = max(1, self.target - self.overlap)
        for i in range(0, len(sentences), step):
            group = sentences[i: i + self.target]
            if not group:
                break
            chunks.append({
                "text": " ".join(group),
                "source": source,
                "chunk_idx": len(chunks),
                "sentence_start": i,
            })
        return chunks


# ─────────────────────────────────────────────
# SMART (semantic boundary detection)
# ─────────────────────────────────────────────

class SmartChunker(BaseChunker):
    """
    Semantic-boundary chunker — splits at:
      ● Markdown headings (#, ##, ###)
      ● Double newlines (paragraph breaks)
      ● Bullet/numbered list boundaries
      ● Long sentence endings after ~min_words words
    
    Merges small fragments up to max_words.
    Result: chunks that respect document structure.
    """

    def __init__(self, min_words: int = 50, max_words: int = 400, overlap_words: int = 30):
        self.min_words = min_words
        self.max_words = max_words
        self.overlap_words = overlap_words

        # Boundary patterns (ordered by priority)
        self._hard_splits = re.compile(
            r"^#{1,6}\s+.+$|^[-=]{3,}$",
            re.MULTILINE
        )
        self._soft_splits = re.compile(
            r"\n{2,}|(?<=\.)\n(?=[A-Z\-\*])"
        )

    def chunk(self, text: str, source: Optional[str] = None) -> List[Dict[str, Any]]:
        # Phase 1: Hard splits
        hard_parts = self._hard_split(text)
        
        # Phase 2: Soft split each part
        raw_chunks = []
        for part in hard_parts:
            raw_chunks.extend(self._soft_split(part))

        # Phase 3: Merge too-small chunks, split too-large
        merged = self._merge_and_split(raw_chunks)

        # Phase 4: Add overlap context
        result = []
        for i, chunk_text in enumerate(merged):
            entry = {
                "text": chunk_text.strip(),
                "source": source,
                "chunk_idx": i,
            }
            if entry["text"]:
                result.append(entry)

        return result

    def _hard_split(self, text: str) -> List[str]:
        parts = self._hard_splits.split(text)
        # Also capture the headings themselves
        headings = self._hard_splits.findall(text)
        result = []
        for i, part in enumerate(parts):
            if i > 0 and headings:
                h = headings[i - 1] if i - 1 < len(headings) else ""
                result.append(h + "\n" + part)
            else:
                result.append(part)
        return [r for r in result if r.strip()]

    def _soft_split(self, text: str) -> List[str]:
        parts = self._soft_splits.split(text)
        return [p.strip() for p in parts if p.strip()]

    def _word_count(self, text: str) -> int:
        return len(text.split())

    def _merge_and_split(self, parts: List[str]) -> List[str]:
        result = []
        buffer = ""

        for part in parts:
            wc_part = self._word_count(part)
            wc_buf = self._word_count(buffer)

            if wc_buf + wc_part <= self.max_words:
                buffer = (buffer + " " + part).strip()
            else:
                if buffer:
                    result.extend(self._force_split(buffer))
                buffer = part

        if buffer:
            result.extend(self._force_split(buffer))

        return result

    def _force_split(self, text: str) -> List[str]:
        """Force-split oversized chunk by word count with overlap."""
        words = text.split()
        if len(words) <= self.max_words:
            return [text]
        
        parts = []
        step = self.max_words - self.overlap_words
        for i in range(0, len(words), step):
            chunk = " ".join(words[i: i + self.max_words])
            if chunk:
                parts.append(chunk)
        return parts


# ─────────────────────────────────────────────
# FACTORY
# ─────────────────────────────────────────────

def AutoChunker(strategy: str = "smart", **kwargs) -> BaseChunker:
    strategies = {
        "smart": SmartChunker,
        "sentence": SentenceChunker,
        "fixed": FixedChunker,
    }
    cls = strategies.get(strategy, SmartChunker)
    return cls(**kwargs)
