from __future__ import annotations

import re
from typing import Callable, Optional

from contexthub.tokenizer import get_tokenizer


def chunk_by_tokens(
    text: str,
    max_tokens: int,
    overlap: int = 0,
    tokenizer: Optional[Callable[[str], int]] = None,
) -> list[str]:
    if not text:
        return []
    if max_tokens <= 0:
        raise ValueError("max_tokens must be > 0")
    if overlap < 0:
        raise ValueError("overlap must be >= 0")
    if overlap >= max_tokens:
        raise ValueError("overlap must be less than max_tokens")

    count = tokenizer or get_tokenizer()
    words = text.split()
    if not words:
        return []

    chunks: list[str] = []
    i = 0
    while i < len(words):
        current_words: list[str] = []
        j = i
        while j < len(words):
            candidate = " ".join(current_words + [words[j]])
            if count(candidate) <= max_tokens or not current_words:
                current_words.append(words[j])
                j += 1
            else:
                break

        chunk = " ".join(current_words)
        if chunk:
            chunks.append(chunk)

        if j >= len(words):
            break

        if overlap == 0:
            i = j
            continue

        overlap_words: list[str] = []
        k = len(current_words) - 1
        while k >= 0:
            candidate = " ".join([current_words[k]] + overlap_words)
            if count(candidate) <= overlap:
                overlap_words.insert(0, current_words[k])
                k -= 1
            else:
                break

        i = max(i + 1, j - len(overlap_words))

    return chunks


def chunk_by_sentences(text: str, max_sentences: int = 5) -> list[str]:
    if not text:
        return []
    if max_sentences <= 0:
        raise ValueError("max_sentences must be > 0")

    sentences = [
        s.strip() for s in re.split(r"(?<=[.!?])\s+", text.strip()) if s.strip()
    ]
    if not sentences:
        return []

    return [
        " ".join(sentences[i : i + max_sentences])
        for i in range(0, len(sentences), max_sentences)
    ]


def chunk_by_paragraphs(text: str) -> list[str]:
    if not text:
        return []
    return [p.strip() for p in text.split("\n\n") if p.strip()]


def chunk_by_separator(text: str, separator: str = "\n") -> list[str]:
    if not text:
        return []
    if separator == "":
        raise ValueError("separator cannot be empty")
    return [part.strip() for part in text.split(separator) if part.strip()]
