import numpy as np

from sentence_api.processing import MOSS_ALGORITHM_VERSION, document_from_transcript
from sentence_api.transcription import (
    Transcript,
    TranscriptionSegment,
    WordTimestamp,
    _wav_bytes,
)


SAMPLE_RATE = 16000
VIDEO_HASH = "a" * 64


def _speech(seconds: float) -> np.ndarray:
    t = np.arange(int(seconds * SAMPLE_RATE)) / SAMPLE_RATE
    return 0.25 * np.sin(2 * np.pi * 220 * t)


def _silence(seconds: float) -> np.ndarray:
    return np.zeros(int(seconds * SAMPLE_RATE))


def test_document_from_transcript_splits_sentences_at_periods(tmp_path):
    samples = np.concatenate(
        [
            _speech(2.0),
            _silence(0.2),
            _speech(2.0),
            _silence(0.2),
            _speech(1.6),
        ]
    ).astype(np.float32)
    wav = tmp_path / "audio.wav"
    wav.write_bytes(_wav_bytes(samples, SAMPLE_RATE))
    transcript = Transcript(
        text="Hello world. Good day everyone. Nice to meet you.",
        segments=[
            TranscriptionSegment(
                start_seconds=0.0,
                end_seconds=2.0,
                text="Hello world.",
                words=[
                    WordTimestamp(0.0, 0.6, "Hello"),
                    WordTimestamp(0.7, 1.5, "world."),
                ],
            ),
            TranscriptionSegment(
                start_seconds=2.2,
                end_seconds=5.5,
                text="Good day everyone. Nice to meet you.",
                words=[
                    WordTimestamp(2.2, 2.8, "Good"),
                    WordTimestamp(2.9, 3.5, "day"),
                    WordTimestamp(3.6, 4.2, "everyone."),
                    WordTimestamp(4.4, 4.9, "Nice"),
                    WordTimestamp(5.0, 5.5, "you."),
                ],
            ),
        ],
    )

    document = document_from_transcript(
        video_hash=VIDEO_HASH,
        duration_ms=6000,
        transcript=transcript,
        language="en",
        audio_path=wav,
    )

    assert document.algorithm_version == MOSS_ALGORITHM_VERSION
    assert [sentence.text for sentence in document.sentences] == [
        "Hello world.",
        "Good day everyone.",
        "Nice to meet you.",
    ]
    assert document.sentences[0].start_ms == 0
    assert document.sentences[0].end_ms == 1500
    assert document.sentences[1].start_ms == 2200
    assert document.sentences[1].end_ms == 4200
    assert document.sentences[2].start_ms == 4200
    assert document.sentences[2].end_ms == 5500
    assert all(sentence.reference_speech_duration_ms > 0 for sentence in document.sentences)