"""Tests for rbsp.core.chunker -- Document Chunking Pipeline."""

from __future__ import annotations

import os
import random
import sys
import time
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from rbsp.core.chunker import (
    STRATEGY_FIXED,
    STRATEGY_PARAGRAPH,
    STRATEGY_SENTENCE,
    Chunk,
    ChunkResult,
    DocumentChunker,
    _count_tokens,
    _token_boundaries,
)


# ---------------------------------------------------------------------------
# Token helpers
# ---------------------------------------------------------------------------


class TestCountTokens(unittest.TestCase):
    """_count_tokens counts alphanumeric/underscore tokens."""

    def test_simple_text(self):
        self.assertEqual(_count_tokens("hello world"), 2)

    def test_with_punctuation(self):
        # Punctuation is not counted as part of tokens
        self.assertEqual(_count_tokens("hello, world!"), 2)

    def test_with_underscores(self):
        self.assertEqual(_count_tokens("foo_bar baz"), 2)

    def test_empty_string(self):
        self.assertEqual(_count_tokens(""), 0)

    def test_only_punctuation(self):
        self.assertEqual(_count_tokens("!@#$%^&*()"), 0)

    def test_mixed(self):
        self.assertEqual(_count_tokens("hello-world foo_bar 123"), 4)


class TestTokenBoundaries(unittest.TestCase):
    """_token_boundaries returns character offset pairs."""

    def test_simple(self):
        result = _token_boundaries("hello world")
        self.assertEqual(result, [(0, 5), (6, 11)])

    def test_empty(self):
        self.assertEqual(_token_boundaries(""), [])

    def test_leading_spaces(self):
        result = _token_boundaries("  hello")
        self.assertEqual(result, [(2, 7)])

    def test_punctuation_gaps(self):
        result = _token_boundaries("a, b! c")
        self.assertEqual(result, [(0, 1), (3, 4), (6, 7)])


# ---------------------------------------------------------------------------
# Chunk and ChunkResult data structures
# ---------------------------------------------------------------------------


class TestChunkDataclass(unittest.TestCase):
    """Chunk is a frozen, slots-based dataclass."""

    def test_creation(self):
        c = Chunk(
            text="hello world", doc_id="doc1", chunk_index=0,
            start_offset=0, end_offset=11, token_count=2,
        )
        self.assertEqual(c.text, "hello world")
        self.assertEqual(c.doc_id, "doc1")
        self.assertEqual(c.chunk_index, 0)
        self.assertEqual(c.start_offset, 0)
        self.assertEqual(c.end_offset, 11)
        self.assertEqual(c.token_count, 2)
        self.assertEqual(c.metadata, {})

    def test_frozen(self):
        c = Chunk(
            text="hi", doc_id="d", chunk_index=0,
            start_offset=0, end_offset=2, token_count=1,
        )
        with self.assertRaises(AttributeError):
            c.text = "bye"

    def test_metadata(self):
        c = Chunk(
            text="hi", doc_id="d", chunk_index=0,
            start_offset=0, end_offset=2, token_count=1,
            metadata={"source": "test"},
        )
        self.assertEqual(c.metadata["source"], "test")


class TestChunkResultDataclass(unittest.TestCase):
    """ChunkResult is a frozen, slots-based dataclass."""

    def test_creation(self):
        c = Chunk(
            text="hello", doc_id="d", chunk_index=0,
            start_offset=0, end_offset=5, token_count=1,
        )
        r = ChunkResult(doc_id="d", chunks=(c,), total_tokens=1, strategy="fixed")
        self.assertEqual(r.doc_id, "d")
        self.assertEqual(len(r.chunks), 1)
        self.assertEqual(r.total_tokens, 1)
        self.assertEqual(r.strategy, "fixed")

    def test_frozen(self):
        r = ChunkResult(doc_id="d", chunks=(), total_tokens=0, strategy="fixed")
        with self.assertRaises(AttributeError):
            r.doc_id = "other"


# ---------------------------------------------------------------------------
# DocumentChunker -- validation
# ---------------------------------------------------------------------------


class TestDocumentChunkerValidation(unittest.TestCase):
    """DocumentChunker constructor validation."""

    def test_default_values(self):
        c = DocumentChunker()
        self.assertIsNotNone(c)

    def test_invalid_chunk_size(self):
        with self.assertRaises(ValueError):
            DocumentChunker(chunk_size=0)

    def test_negative_overlap(self):
        with self.assertRaises(ValueError):
            DocumentChunker(overlap=-1)

    def test_overlap_ge_chunk_size(self):
        with self.assertRaises(ValueError):
            DocumentChunker(chunk_size=10, overlap=10)

    def test_invalid_strategy(self):
        with self.assertRaises(ValueError):
            DocumentChunker(strategy="invalid")

    def test_invalid_min_chunk_size(self):
        with self.assertRaises(ValueError):
            DocumentChunker(min_chunk_size=0)

    def test_valid_strategies(self):
        for s in (STRATEGY_FIXED, STRATEGY_SENTENCE, STRATEGY_PARAGRAPH):
            c = DocumentChunker(strategy=s)
            self.assertIsNotNone(c)

    def test_repr(self):
        c = DocumentChunker(chunk_size=100, overlap=10, strategy="fixed")
        r = repr(c)
        self.assertIn("DocumentChunker", r)
        self.assertIn("chunk_size=100", r)
        self.assertIn("overlap=10", r)
        self.assertIn("fixed", r)


# ---------------------------------------------------------------------------
# DocumentChunker -- empty / edge cases
# ---------------------------------------------------------------------------


class TestDocumentChunkerEdgeCases(unittest.TestCase):
    """Edge cases: empty text, whitespace, single token."""

    def test_empty_string(self):
        chunker = DocumentChunker(strategy="fixed")
        result = chunker.chunk("")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.chunks[0].text, "")
        self.assertEqual(result.total_tokens, 0)

    def test_whitespace_only(self):
        chunker = DocumentChunker(strategy="fixed")
        result = chunker.chunk("   \t\n  ")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.chunks[0].text, "")
        self.assertEqual(result.total_tokens, 0)

    def test_single_token(self):
        chunker = DocumentChunker(chunk_size=10, overlap=0, strategy="fixed")
        result = chunker.chunk("hello")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.chunks[0].text, "hello")
        self.assertEqual(result.chunks[0].token_count, 1)

    def test_doc_id_propagated(self):
        chunker = DocumentChunker(strategy="fixed")
        result = chunker.chunk("hello world", doc_id="my_doc.txt")
        self.assertEqual(result.doc_id, "my_doc.txt")
        for c in result.chunks:
            self.assertEqual(c.doc_id, "my_doc.txt")

    def test_empty_doc_id_default(self):
        chunker = DocumentChunker(strategy="fixed")
        result = chunker.chunk("hello world")
        self.assertEqual(result.doc_id, "")

    def test_strategy_in_result(self):
        for s in ("fixed", "sentence", "paragraph"):
            chunker = DocumentChunker(strategy=s, chunk_size=1000)
            result = chunker.chunk("Hello world.")
            self.assertEqual(result.strategy, s)


# ---------------------------------------------------------------------------
# DocumentChunker -- fixed strategy
# ---------------------------------------------------------------------------


class TestFixedStrategy(unittest.TestCase):
    """Fixed-window chunking strategy."""

    def test_short_text_single_chunk(self):
        chunker = DocumentChunker(chunk_size=10, overlap=0, strategy="fixed")
        result = chunker.chunk("hello world")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.chunks[0].token_count, 2)

    def test_exact_chunk_size(self):
        chunker = DocumentChunker(chunk_size=5, overlap=0, strategy="fixed")
        result = chunker.chunk("one two three four five")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.chunks[0].token_count, 5)

    def test_multiple_chunks_no_overlap(self):
        chunker = DocumentChunker(chunk_size=3, overlap=0, strategy="fixed", min_chunk_size=1)
        result = chunker.chunk("a b c d e f")
        self.assertEqual(len(result.chunks), 2)
        self.assertEqual(result.chunks[0].token_count, 3)
        self.assertEqual(result.chunks[1].token_count, 3)

    def test_overlap(self):
        chunker = DocumentChunker(chunk_size=3, overlap=1, strategy="fixed", min_chunk_size=1)
        result = chunker.chunk("a b c d e")
        # step=2: [0:3]="a b c" (3 tok), [2:5]="c d e" (3 tok), [4:5]="e" (1 tok)
        self.assertGreaterEqual(len(result.chunks), 2)

    def test_offsets_correct(self):
        chunker = DocumentChunker(chunk_size=2, overlap=0, strategy="fixed", min_chunk_size=1)
        result = chunker.chunk("hello world foo bar")
        c0 = result.chunks[0]
        self.assertEqual(c0.start_offset, 0)
        self.assertEqual(c0.end_offset, 11)
        self.assertEqual(c0.text, "hello world")

    def test_chunk_indices_sequential(self):
        chunker = DocumentChunker(chunk_size=2, overlap=0, strategy="fixed", min_chunk_size=1)
        result = chunker.chunk("a b c d e f")
        for i, c in enumerate(result.chunks):
            self.assertEqual(c.chunk_index, i)

    def test_total_tokens(self):
        chunker = DocumentChunker(chunk_size=3, overlap=0, strategy="fixed", min_chunk_size=1)
        text = "one two three four five six"
        result = chunker.chunk(text)
        # Total tokens should cover all tokens (possibly with overlap)
        self.assertGreaterEqual(result.total_tokens, 6)


# ---------------------------------------------------------------------------
# DocumentChunker -- sentence strategy
# ---------------------------------------------------------------------------


class TestSentenceStrategy(unittest.TestCase):
    """Sentence-boundary chunking strategy."""

    def test_single_sentence(self):
        chunker = DocumentChunker(chunk_size=100, strategy="sentence")
        result = chunker.chunk("Hello world this is a test.")
        self.assertEqual(len(result.chunks), 1)

    def test_multiple_sentences(self):
        text = "First sentence. Second sentence. Third sentence."
        chunker = DocumentChunker(chunk_size=3, overlap=0, strategy="sentence", min_chunk_size=1)
        result = chunker.chunk(text)
        self.assertGreater(len(result.chunks), 1)

    def test_respects_sentence_boundaries(self):
        text = "Hello world. Foo bar baz. Another one."
        chunker = DocumentChunker(chunk_size=3, overlap=0, strategy="sentence", min_chunk_size=1)
        result = chunker.chunk(text)
        # Each chunk should end at a sentence boundary (period)
        for c in result.chunks:
            stripped = c.text.strip()
            if stripped:
                # The last char or nearby should be a period (allowing trailing)
                self.assertTrue(
                    stripped.endswith(".") or stripped.endswith("!") or stripped.endswith("?")
                    or _count_tokens(stripped) <= chunker._chunk_size,
                    f"Chunk doesn't respect boundary: {stripped!r}"
                )

    def test_question_and_exclamation(self):
        text = "Is this a test? Yes it is! Great."
        chunker = DocumentChunker(chunk_size=100, strategy="sentence")
        result = chunker.chunk(text)
        self.assertEqual(len(result.chunks), 1)

    def test_default_strategy_is_sentence(self):
        chunker = DocumentChunker()
        result = chunker.chunk("Hello. World.")
        self.assertEqual(result.strategy, "sentence")


# ---------------------------------------------------------------------------
# DocumentChunker -- paragraph strategy
# ---------------------------------------------------------------------------


class TestParagraphStrategy(unittest.TestCase):
    """Paragraph-boundary chunking strategy."""

    def test_single_paragraph(self):
        chunker = DocumentChunker(chunk_size=100, strategy="paragraph")
        result = chunker.chunk("This is a single paragraph with no blank lines.")
        self.assertEqual(len(result.chunks), 1)

    def test_multiple_paragraphs(self):
        text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
        chunker = DocumentChunker(chunk_size=3, overlap=0, strategy="paragraph", min_chunk_size=1)
        result = chunker.chunk(text)
        self.assertGreater(len(result.chunks), 1)

    def test_blank_line_separation(self):
        text = "Para one.\n\nPara two.\n\nPara three."
        chunker = DocumentChunker(chunk_size=100, strategy="paragraph")
        result = chunker.chunk(text)
        # All paragraphs fit in one chunk
        self.assertEqual(len(result.chunks), 1)

    def test_large_paragraphs_split(self):
        rng = random.Random(42)
        para1 = " ".join(f"word{rng.randint(0,99)}" for _ in range(50))
        para2 = " ".join(f"word{rng.randint(0,99)}" for _ in range(50))
        text = f"{para1}\n\n{para2}"
        chunker = DocumentChunker(chunk_size=30, overlap=0, strategy="paragraph", min_chunk_size=1)
        result = chunker.chunk(text)
        self.assertGreater(len(result.chunks), 1)

    def test_multiple_blank_lines(self):
        text = "One.\n\n\n\nTwo."
        chunker = DocumentChunker(chunk_size=100, strategy="paragraph")
        result = chunker.chunk(text)
        self.assertGreater(result.total_tokens, 0)


# ---------------------------------------------------------------------------
# DocumentChunker -- min_chunk_size merging
# ---------------------------------------------------------------------------


class TestMinChunkMerge(unittest.TestCase):
    """Trailing chunks below min_chunk_size merge into previous."""

    def test_tiny_trailing_merged(self):
        # 6 tokens, chunk_size=5, overlap=0 -> chunks of [5, 1]
        # min_chunk_size=2 -> the trailing 1-token chunk merges into previous
        chunker = DocumentChunker(
            chunk_size=5, overlap=0, strategy="fixed", min_chunk_size=2
        )
        result = chunker.chunk("a b c d e f")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.total_tokens, 6)

    def test_no_merge_when_above_min(self):
        # 8 tokens, chunk_size=5, overlap=0 -> chunks of [5, 3]
        # min_chunk_size=2 -> 3 >= 2, no merge
        chunker = DocumentChunker(
            chunk_size=5, overlap=0, strategy="fixed", min_chunk_size=2
        )
        result = chunker.chunk("a b c d e f g h")
        self.assertEqual(len(result.chunks), 2)

    def test_single_chunk_never_merged(self):
        # Even if single chunk has fewer tokens than min_chunk_size, no merge needed
        chunker = DocumentChunker(
            chunk_size=100, overlap=0, strategy="fixed", min_chunk_size=50
        )
        result = chunker.chunk("a b c")
        self.assertEqual(len(result.chunks), 1)
        self.assertEqual(result.chunks[0].token_count, 3)


# ---------------------------------------------------------------------------
# DocumentChunker -- chunk_many
# ---------------------------------------------------------------------------


class TestChunkMany(unittest.TestCase):
    """chunk_many processes multiple documents."""

    def test_basic(self):
        chunker = DocumentChunker(chunk_size=100, strategy="fixed")
        docs = [
            ("doc1", "Hello world."),
            ("doc2", "Foo bar baz."),
        ]
        results = chunker.chunk_many(docs)
        self.assertEqual(len(results), 2)
        self.assertEqual(results[0].doc_id, "doc1")
        self.assertEqual(results[1].doc_id, "doc2")

    def test_empty_list(self):
        chunker = DocumentChunker(strategy="fixed")
        results = chunker.chunk_many([])
        self.assertEqual(results, [])

    def test_doc_ids_propagated(self):
        chunker = DocumentChunker(chunk_size=100, strategy="fixed")
        docs = [(f"doc_{i}", f"Content {i}") for i in range(5)]
        results = chunker.chunk_many(docs)
        for i, r in enumerate(results):
            self.assertEqual(r.doc_id, f"doc_{i}")


# ---------------------------------------------------------------------------
# Unicode and special text
# ---------------------------------------------------------------------------


class TestUnicodeChunking(unittest.TestCase):
    """Unicode and special characters in chunking."""

    def test_unicode_text(self):
        chunker = DocumentChunker(chunk_size=100, strategy="fixed")
        result = chunker.chunk("Caf\u00e9 \u00fc\u00f1\u00eec\u00f6d\u00e9 \u00e4\u00f6\u00fc\u00df")
        self.assertGreater(result.total_tokens, 0)

    def test_mixed_scripts(self):
        chunker = DocumentChunker(chunk_size=100, strategy="fixed")
        result = chunker.chunk("hello world 123 test")
        self.assertEqual(result.total_tokens, 4)

    def test_code_text(self):
        chunker = DocumentChunker(chunk_size=5, overlap=0, strategy="fixed", min_chunk_size=1)
        code = "def foo_bar(x): return x + 1"
        result = chunker.chunk(code)
        self.assertGreater(result.total_tokens, 0)


# ---------------------------------------------------------------------------
# Performance
# ---------------------------------------------------------------------------


class TestChunkerPerformance(unittest.TestCase):
    """Chunking performance benchmarks."""

    def test_10000_tokens_under_100ms(self):
        words = [f"word{i}" for i in range(10000)]
        text = " ".join(words)
        chunker = DocumentChunker(chunk_size=384, overlap=64, strategy="fixed")

        t0 = time.perf_counter()
        result = chunker.chunk(text)
        elapsed = time.perf_counter() - t0

        self.assertGreater(len(result.chunks), 1)
        self.assertLess(elapsed, 2.0, f"10k tokens took {elapsed:.3f}s")

    def test_sentence_strategy_performance(self):
        rng = random.Random(42)
        sentences = [
            " ".join(f"word{rng.randint(0,999)}" for _ in range(10)) + "."
            for _ in range(500)
        ]
        text = " ".join(sentences)
        chunker = DocumentChunker(chunk_size=100, overlap=10, strategy="sentence")

        t0 = time.perf_counter()
        result = chunker.chunk(text)
        elapsed = time.perf_counter() - t0

        self.assertGreater(len(result.chunks), 1)
        self.assertLess(elapsed, 1.0, f"Sentence chunking took {elapsed:.3f}s")


if __name__ == "__main__":
    unittest.main()
