from __future__ import annotations

import unittest

from cognitive_agent.infrastructure.speech.assembler import TranscriptAssembler
from cognitive_agent.infrastructure.speech.contracts import SpeechSegmentResult


def segment(
    segment_id: str,
    sequence: int,
    text: str,
    *,
    final: bool,
    start_ms: int | None = None,
    end_ms: int | None = None,
    overlap_ms: int = 0,
) -> SpeechSegmentResult:
    start = sequence * 1000 if start_ms is None else start_ms
    end = start + 1200 if end_ms is None else end_ms
    return SpeechSegmentResult(
        utterance_id="utterance-1",
        segment_id=segment_id,
        sequence=sequence,
        start_ms=start,
        end_ms=end,
        text=text,
        is_final=final,
        audio_overlap_ms=overlap_ms,
    )


class TranscriptAssemblerTests(unittest.TestCase):
    def test_partial_replaces_partial_and_final_commits(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(segment("s0", 0, "打开空", final=False))
        self.assertEqual(assembler.snapshot().text, "打开空")
        assembler.update(segment("s0", 0, "打开空调", final=False))
        self.assertEqual(assembler.snapshot().text, "打开空调")
        final = assembler.update(segment("s0", 0, "打开空调。", final=True))
        self.assertEqual(final.text, "打开空调。")
        self.assertFalse(final.partial_segments)

    def test_duplicate_final_callback_is_idempotent(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        callback = segment("s0", 0, "打开空调。", final=True)
        assembler.update(callback)
        result = assembler.update(callback)
        self.assertEqual(result.text, "打开空调。")
        self.assertEqual(result.duplicate_callbacks, 1)

    def test_final_cannot_be_mutated(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(segment("s0", 0, "打开空调。", final=True))
        result = assembler.update(segment("s0", 0, "关闭空调。", final=True))
        self.assertTrue(result.uncertain)
        self.assertEqual(result.text, "打开空调。")

    def test_out_of_order_callbacks_are_sorted(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(segment("s1", 1, "打开氛围灯。", final=True))
        result = assembler.update(segment("s0", 0, "打开空调。", final=True))
        self.assertEqual(result.text, "打开空调。打开氛围灯。")
        self.assertEqual(result.out_of_order_callbacks, 1)

    def test_audio_supported_boundary_overlap_is_removed(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(
            segment("s0", 0, "把座椅靠背", final=True, start_ms=0, end_ms=1500)
        )
        result = assembler.update(
            segment(
                "s1",
                1,
                "靠背向后调24度",
                final=True,
                start_ms=1400,
                end_ms=2600,
                overlap_ms=100,
            )
        )
        self.assertEqual(result.text, "把座椅靠背，向后调24度")
        self.assertEqual(result.overlaps_removed, ("靠背",))

    def test_long_text_overlap_is_removed_without_audio_overlap(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(segment("s0", 0, "请打开空调温度", final=True))
        result = assembler.update(segment("s1", 1, "空调温度调到24度", final=True))
        self.assertIn("空调温度", result.overlaps_removed)

    def test_repeated_action_in_distinct_commands_is_preserved(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(segment("s0", 0, "打开前排空调", final=True))
        result = assembler.update(segment("s1", 1, "打开后排空调", final=True))
        self.assertEqual(result.text, "打开前排空调，打开后排空调")
        self.assertFalse(result.overlaps_removed)

    def test_finish_rejects_late_callbacks_and_reset_starts_new_utterance(self) -> None:
        assembler = TranscriptAssembler("utterance-1")
        assembler.update(segment("s0", 0, "打开空调", final=True))
        assembler.finish()
        late = assembler.update(segment("s1", 1, "关闭空调", final=True))
        self.assertEqual(late.text, "打开空调")
        assembler.reset("utterance-2")
        self.assertEqual(assembler.snapshot().utterance_id, "utterance-2")
        self.assertEqual(assembler.snapshot().text, "")


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