"""Tests covering previously untested behaviours in chunk_sentences."""

import json
from pathlib import Path

from pymultirole_plugins.v1.schema import Document, Sentence

from pyprocessors_chunk_sentences.chunk_sentences import (
    CharacterBasedSentenceAnalysis,
    ChunkingUnit,
    ChunkSentencesParameters,
    ChunkSentencesProcessor,
    MarkedSentence,
    ResegmenterParameters,
    ResegmenterProcessor,
    SentenceResegmenter,
    TikTokenBasedSentenceAnalysis,
    TokenModel,
    WrappedSentence,
    analyze_sentence,
    get_model,
    get_tiktoken_encoding,
)

# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------


def _load_document(name: str) -> Document:
    testdir = Path(__file__).parent
    source = Path(testdir, f"data/{name}")
    with source.open("r") as f:
        return Document(**json.load(f))


# ---------------------------------------------------------------------------
# MarkedSentence
# ---------------------------------------------------------------------------


class TestMarkedSentence:
    def test_properties_read(self):
        s = Sentence(start=10, end=20, metadata={"key": "val"})
        ms = MarkedSentence(sentence=s, is_marked=True, length=10)
        assert ms.start == 10
        assert ms.end == 20
        assert ms.metadata == {"key": "val"}
        assert ms.is_marked is True
        assert ms.length == 10

    def test_properties_write(self):
        s = Sentence(start=0, end=5)
        ms = MarkedSentence(sentence=s)
        ms.start = 3
        ms.end = 8
        assert ms.start == 3
        assert ms.end == 8
        # Underlying sentence is mutated
        assert s.start == 3
        assert s.end == 8

    def test_defaults(self):
        s = Sentence(start=0, end=1)
        ms = MarkedSentence(sentence=s)
        assert ms.is_marked is False
        assert ms.length == -1


# ---------------------------------------------------------------------------
# WrappedSentence
# ---------------------------------------------------------------------------


class TestWrappedSentence:
    def test_defaults_from_marked_sentence(self):
        s = Sentence(start=5, end=15)
        ms = MarkedSentence(sentence=s, length=10)
        ws = WrappedSentence(sentence=ms)
        assert ws.wrapped is ms
        assert ws.start == 5
        assert ws.end == 15

    def test_custom_start_end(self):
        s = Sentence(start=5, end=15)
        ms = MarkedSentence(sentence=s, length=10)
        ws = WrappedSentence(sentence=ms, start=0, end=10)
        assert ws.start == 0
        assert ws.end == 10


# ---------------------------------------------------------------------------
# CharacterBasedSentenceAnalysis
# ---------------------------------------------------------------------------


class TestCharacterBasedSentenceAnalysis:
    def test_len(self):
        text = "Hello, world!"
        s = Sentence(start=0, end=len(text))
        analysis = CharacterBasedSentenceAnalysis(text, s)
        analysis.compute()
        assert len(analysis) == 13

    def test_split_basic(self):
        text = "abcdefghij"  # 10 chars
        s = Sentence(start=0, end=10)
        analysis = CharacterBasedSentenceAnalysis(text, s)
        analysis.compute()
        parts = analysis.split(4)
        # 10 chars split at max_len=4 → 3 parts: [0,4), [4,8), [8,10)
        assert len(parts) == 3
        assert parts[0].sentence.start == 0
        assert parts[0].sentence.end == 4
        assert parts[0].length == 4
        assert parts[1].sentence.start == 4
        assert parts[1].sentence.end == 8
        assert parts[2].sentence.start == 8
        assert parts[2].sentence.end == 10
        assert parts[2].length == 2

    def test_split_with_whitespace_stripping(self):
        text = "  hello  "
        s = Sentence(start=0, end=9)
        analysis = CharacterBasedSentenceAnalysis(text, s)
        analysis.compute()
        parts = analysis.split(20)
        # Single part, whitespace stripped
        assert len(parts) == 1
        assert parts[0].sentence.start == 2
        assert parts[0].sentence.end == 7

    def test_split_preserves_metadata(self):
        text = "abcdefghij"
        s = Sentence(start=0, end=10, metadata={"tag": "test"})
        analysis = CharacterBasedSentenceAnalysis(text, s)
        analysis.compute()
        parts = analysis.split(5)
        for part in parts:
            assert part.sentence.metadata == {"tag": "test"}

    def test_maybe_split_no_split_needed(self):
        text = "short"
        s = Sentence(start=0, end=5)
        analysis = CharacterBasedSentenceAnalysis(text, s)
        analysis.compute()
        parts = analysis.maybe_split(100)
        assert len(parts) == 1
        assert parts[0].sentence is s
        assert parts[0].length == 5

    def test_maybe_split_split_needed(self):
        text = "abcdefghij"
        s = Sentence(start=0, end=10)
        analysis = CharacterBasedSentenceAnalysis(text, s)
        analysis.compute()
        parts = analysis.maybe_split(4)
        assert len(parts) == 3


# ---------------------------------------------------------------------------
# TikTokenBasedSentenceAnalysis
# ---------------------------------------------------------------------------


class TestTikTokenBasedSentenceAnalysis:
    def test_encode_decode(self):
        text = "Hello, world!"
        s = Sentence(start=0, end=len(text))
        analysis = TikTokenBasedSentenceAnalysis(text, s, "gpt-4")
        analysis.compute()
        assert len(analysis) > 0
        # Round-trip
        tokens = analysis.encode(text)
        decoded = analysis.decode(tokens)
        assert decoded == text

    def test_split_round_trip(self):
        text = "The quick brown fox jumps over the lazy dog. " * 5
        s = Sentence(start=0, end=len(text))
        analysis = TikTokenBasedSentenceAnalysis(text, s, "gpt-4")
        analysis.compute()
        parts = analysis.split(5)
        assert len(parts) > 1
        # All parts should have valid offsets within the text
        for part in parts:
            assert 0 <= part.sentence.start <= part.sentence.end <= len(text)
            sub_text = text[part.sentence.start : part.sentence.end]
            assert len(sub_text) > 0


# ---------------------------------------------------------------------------
# get_tiktoken_encoding (caching)
# ---------------------------------------------------------------------------


class TestGetTiktokenEncoding:
    def test_returns_encoding(self):
        enc = get_tiktoken_encoding("gpt-4")
        assert enc is not None
        tokens = enc.encode("hello")
        assert len(tokens) > 0

    def test_caching(self):
        enc1 = get_tiktoken_encoding("gpt-4")
        enc2 = get_tiktoken_encoding("gpt-4")
        assert enc1 is enc2


# ---------------------------------------------------------------------------
# analyze_sentence
# ---------------------------------------------------------------------------


class TestAnalyzeSentence:
    def test_character_unit(self):
        text = "Hello world"
        s = Sentence(start=0, end=len(text))
        params = ChunkSentencesParameters(unit=ChunkingUnit.character)
        analysis = analyze_sentence(params, text, s)
        assert isinstance(analysis, CharacterBasedSentenceAnalysis)
        assert len(analysis) == len(text)

    def test_token_unit_gpt(self):
        text = "Hello world"
        s = Sentence(start=0, end=len(text))
        params = ChunkSentencesParameters(unit=ChunkingUnit.token, model=TokenModel.gpt_4)
        analysis = analyze_sentence(params, text, s)
        assert isinstance(analysis, TikTokenBasedSentenceAnalysis)
        assert len(analysis) > 0

    def test_token_unit_text_embedding(self):
        text = "Hello world"
        s = Sentence(start=0, end=len(text))
        params = ChunkSentencesParameters(unit=ChunkingUnit.token, model=TokenModel.text_embedding_3_large)
        analysis = analyze_sentence(params, text, s)
        assert isinstance(analysis, TikTokenBasedSentenceAnalysis)

    def test_token_unit_fallback_model(self):
        """Non-GPT, non-embedding models should fall back to text-embedding-ada-002."""
        text = "Hello world"
        s = Sentence(start=0, end=len(text))
        params = ChunkSentencesParameters(unit=ChunkingUnit.token, model=TokenModel.wbd)
        analysis = analyze_sentence(params, text, s)
        assert isinstance(analysis, TikTokenBasedSentenceAnalysis)
        # The encoding should be ada-002's encoding
        assert analysis.encoding.name == "cl100k_base"


# ---------------------------------------------------------------------------
# get_model
# ---------------------------------------------------------------------------


class TestGetModel:
    def test_wbd_returns_minus_one(self):
        assert get_model("wbd") == -1

    def test_gpt4_returns_tiktoken_encoding(self):
        h = get_model("gpt-4")
        tokens = h.encode("hello")
        assert len(tokens) > 0

    def test_text_embedding_returns_tiktoken_encoding(self):
        h = get_model("text-embedding-ada-002")
        tokens = h.encode("hello")
        assert len(tokens) > 0


# ---------------------------------------------------------------------------
# ChunkSentencesProcessor — edge cases
# ---------------------------------------------------------------------------


class TestChunkSentencesProcessorEdgeCases:
    def test_empty_sentences(self):
        """Documents with no sentences should be returned unchanged."""
        doc = Document(text="Some text", sentences=[])
        processor = ChunkSentencesProcessor()
        params = ChunkSentencesParameters()
        result = processor.process([doc], params)
        assert len(result) == 1
        assert result[0].sentences == []

    def test_multiple_documents(self):
        """Processor should handle multiple documents."""
        doc1 = Document(
            text="First sentence. Second sentence.",
            sentences=[Sentence(start=0, end=15), Sentence(start=16, end=32)],
        )
        doc2 = Document(
            text="Another text here. More text follows.",
            sentences=[Sentence(start=0, end=18), Sentence(start=19, end=37)],
        )
        processor = ChunkSentencesProcessor()
        params = ChunkSentencesParameters(chunk_char_max_length=5000)
        result = processor.process([doc1, doc2], params)
        assert len(result) == 2
        assert len(result[0].sentences) >= 1
        assert len(result[1].sentences) >= 1

    def test_single_sentence_document(self):
        """A document with exactly one sentence should produce one chunk."""
        text = "Hello world."
        doc = Document(text=text, sentences=[Sentence(start=0, end=len(text))])
        processor = ChunkSentencesProcessor()
        params = ChunkSentencesParameters(chunk_char_max_length=5000)
        result = processor.process([doc], params)
        assert len(result[0].sentences) == 1
        assert result[0].sentences[0].start == 0
        assert result[0].sentences[0].end == len(text)

    def test_chunk_char_splits_long_sentences(self):
        """Sentences longer than chunk_char_max_length should be split."""
        text = "a" * 100
        doc = Document(text=text, sentences=[Sentence(start=0, end=100)])
        processor = ChunkSentencesProcessor()
        params = ChunkSentencesParameters(chunk_char_max_length=30, unit=ChunkingUnit.character)
        result = processor.process([doc], params)
        assert len(result[0].sentences) > 1
        for s in result[0].sentences:
            assert s.end - s.start <= 30

    def test_overlap_one(self):
        """Test overlap=1 produces overlapping chunks."""
        # Create a document with many small sentences
        text = "aa. bb. cc. dd. ee. ff."
        sentences = [
            Sentence(start=0, end=2),
            Sentence(start=4, end=6),
            Sentence(start=8, end=10),
            Sentence(start=12, end=14),
            Sentence(start=16, end=18),
            Sentence(start=20, end=22),
        ]
        doc = Document(text=text, sentences=sentences)
        processor = ChunkSentencesProcessor()
        # Small chunk size to force multiple chunks, overlap=1
        params = ChunkSentencesParameters(chunk_char_max_length=7, overlap=1)
        result = processor.process([doc], params)
        chunks = result[0].sentences
        # With overlap, subsequent chunks should start before the previous one ends
        if len(chunks) > 1:
            for i in range(1, len(chunks)):
                assert chunks[i].start < chunks[i - 1].end

    def test_get_model_returns_parameters_class(self):
        assert ChunkSentencesProcessor.get_model() == ChunkSentencesParameters

    def test_tokenize_with_model_wbd(self):
        """tokenize_with_model with model=-1 should use blingfire.text_to_words."""
        tokens = ChunkSentencesProcessor.tokenize_with_model(-1, "Hello world")
        assert isinstance(tokens, list)
        assert len(tokens) >= 2

    def test_tokenize_with_model_tiktoken(self):
        """tokenize_with_model with a tiktoken encoding should return token ids."""
        h = get_model("gpt-4")
        tokens = ChunkSentencesProcessor.tokenize_with_model(h, "Hello world")
        assert len(tokens) >= 1


# ---------------------------------------------------------------------------
# ResegmenterProcessor
# ---------------------------------------------------------------------------


class TestResegmenterProcessor:
    def test_get_model_returns_parameters_class(self):
        assert ResegmenterProcessor.get_model() == ResegmenterParameters

    def test_resegmenter_with_syntok(self):
        """Resegmenter should work with the syntok segmenter."""
        doc = _load_document("bigsegments_segmented.json")
        processor = ResegmenterProcessor()
        params = ResegmenterParameters(
            unit=ChunkingUnit.token,
            chunk_token_max_length=256,
            sentence_resegmenter=SentenceResegmenter.syntok,
        )
        result = processor.process([doc], params)
        chunked = result[0]
        # Should produce more sentences than the original (resegmented)
        assert len(chunked.sentences) > 0

    def test_resegmenter_no_sentences(self):
        """Resegmenter with no pre-existing sentences should segment the whole text."""
        text = "First sentence. Second sentence. Third sentence."
        doc = Document(text=text, sentences=[])
        processor = ResegmenterProcessor()
        params = ResegmenterParameters(
            unit=ChunkingUnit.character,
            chunk_char_max_length=5000,
        )
        result = processor.process([doc], params)
        chunked = result[0]
        # Should have created sentences from the whole text
        assert len(chunked.sentences) >= 1

    def test_resegmenter_multiple_documents(self):
        """Resegmenter should handle multiple documents."""
        text1 = "Sentence one. Sentence two. Sentence three."
        text2 = "Another text. With multiple sentences. Here."
        doc1 = Document(text=text1, sentences=[Sentence(start=0, end=len(text1))])
        doc2 = Document(text=text2, sentences=[Sentence(start=0, end=len(text2))])
        processor = ResegmenterProcessor()
        params = ResegmenterParameters(
            unit=ChunkingUnit.character,
            chunk_char_max_length=5000,
        )
        result = processor.process([doc1, doc2], params)
        assert len(result) == 2
        assert len(result[0].sentences) >= 1
        assert len(result[1].sentences) >= 1


# ---------------------------------------------------------------------------
# Parameters / model classes
# ---------------------------------------------------------------------------


class TestParameters:
    def test_chunk_sentences_parameters_defaults(self):
        p = ChunkSentencesParameters()
        assert p.unit == ChunkingUnit.character
        assert p.model == TokenModel.wbd
        assert p.chunk_char_max_length == 1024
        assert p.chunk_token_max_length == 128
        assert p.overlap == 0
        assert p.trunc_char_max_length == 0
        assert p.trunc_token_max_length == 0

    def test_resegmenter_parameters_defaults(self):
        p = ResegmenterParameters()
        assert p.sentence_resegmenter == SentenceResegmenter.blingfire

    def test_resegmenter_parameters_syntok(self):
        p = ResegmenterParameters(sentence_resegmenter=SentenceResegmenter.syntok)
        assert p.sentence_resegmenter == SentenceResegmenter.syntok

    def test_chunking_unit_values(self):
        assert ChunkingUnit.character == "character"
        assert ChunkingUnit.token == "token"

    def test_token_model_values(self):
        assert TokenModel.wbd == "wbd"
        assert TokenModel.gpt_4 == "gpt-4"
        assert TokenModel.gpt_4o == "gpt-4o"
        assert TokenModel.text_embedding_ada_002 == "text-embedding-ada-002"
