import pytest

from audita.core.chunking import TokenEstimatorProtocol, chunk_transcript
from audita.core.errors import AuditaValidationError
from audita.core.schemas import parse_transcript_json


class FakeEstimator(TokenEstimatorProtocol):
    def estimate_json(self, value):
        if len(value) == 1:
            return 4
        return len(value) * 4


def test_chunk_transcript_batches_sections_by_token_limit():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
        ]
        """
    )

    sections = chunk_transcript(transcript, max_section_tokens=8, estimator=FakeEstimator())

    assert len(sections) == 2
    assert [segment.segment.id for segment in sections[0].segments] == [1, 2]
    assert [segment.segment.id for segment in sections[1].segments] == [3]


def test_chunk_transcript_targets_llm_concurrency_when_feasible():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
          {"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"}
        ]
        """
    )

    sections = chunk_transcript(
        transcript,
        max_section_tokens=8,
        min_section_tokens=4,
        target_section_count=2,
        estimator=FakeEstimator(),
    )

    assert len(sections) == 2
    assert [segment.segment.id for segment in sections[0].segments] == [1, 2]
    assert [segment.segment.id for segment in sections[1].segments] == [3, 4]


def test_chunk_transcript_increases_section_count_when_target_sections_exceed_max_tokens():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
          {"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"},
          {"id": 5, "speaker": "A", "start": 4.0, "end": 5.0, "text": "five"}
        ]
        """
    )

    sections = chunk_transcript(
        transcript,
        max_section_tokens=8,
        min_section_tokens=4,
        target_section_count=1,
        estimator=FakeEstimator(),
    )

    assert len(sections) > 1
    assert all(section.token_count <= 8 for section in sections)


def test_chunk_transcript_reduces_section_count_when_target_sections_fall_below_min_tokens():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
        ]
        """
    )

    sections = chunk_transcript(
        transcript,
        max_section_tokens=12,
        min_section_tokens=8,
        target_section_count=3,
        estimator=FakeEstimator(),
    )

    assert len(sections) == 1
    assert sections[0].token_count >= 8


def test_chunk_transcript_exact_target_sections_returns_exact_count_when_feasible():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
          {"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"}
        ]
        """
    )

    sections = chunk_transcript(
        transcript,
        max_section_tokens=8,
        min_section_tokens=4,
        exact_target_section_count=2,
        estimator=FakeEstimator(),
    )

    assert len(sections) == 2
    assert [segment.segment.id for segment in sections[0].segments] == [1, 2]
    assert [segment.segment.id for segment in sections[1].segments] == [3, 4]


def test_chunk_transcript_exact_target_sections_errors_when_target_exceeds_segment_count():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"}
        ]
        """
    )

    with pytest.raises(AuditaValidationError, match="Target section count exceeds the number of transcript segments"):
        chunk_transcript(
            transcript,
            max_section_tokens=8,
            min_section_tokens=4,
            exact_target_section_count=3,
            estimator=FakeEstimator(),
        )


def test_chunk_transcript_exact_target_sections_errors_when_section_would_exceed_max():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
        ]
        """
    )

    with pytest.raises(AuditaValidationError, match="AUDITA_TARGET_SECTIONS cannot produce contiguous transcript sections"):
        chunk_transcript(
            transcript,
            max_section_tokens=8,
            min_section_tokens=4,
            exact_target_section_count=1,
            estimator=FakeEstimator(),
        )


def test_chunk_transcript_exact_target_sections_errors_when_section_would_fall_below_min():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
          {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
        ]
        """
    )

    with pytest.raises(AuditaValidationError, match="AUDITA_TARGET_SECTIONS cannot produce contiguous transcript sections"):
        chunk_transcript(
            transcript,
            max_section_tokens=12,
            min_section_tokens=8,
            exact_target_section_count=3,
            estimator=FakeEstimator(),
        )


def test_chunk_transcript_prompt_payload_includes_categories_when_present():
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one", "categories": ["intro", "aside"]}
        ]
        """
    )

    sections = chunk_transcript(transcript, max_section_tokens=8, estimator=FakeEstimator())

    assert sections[0].prompt_payload() == [
        {"id": 1, "original_text": "one", "categories": ["intro", "aside"]}
    ]