import numpy as np

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


SAMPLE_RATE = 16000
VIDEO_HASH = "a" * 64


class WordTimestampTranscriber:
    def __init__(self, transcript):
        self.transcript = transcript

    @property
    def available(self):
        return True

    def transcribe(self, audio_path, language=None):
        return self.transcript


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_punctuation(tmp_path):
    samples = np.concatenate(
        [
            _speech(8.0),
            _speech(1.5),
            _speech(10.5),
            _speech(4.0),
        ]
    ).astype(np.float32)
    wav = tmp_path / "audio.wav"
    wav.write_bytes(_wav_bytes(samples, SAMPLE_RATE))
    transcript = Transcript(
        text=(
            "So I've been single for about four years now and I find it hard to meet "
            "a guy especially in California. I've tried dating apps, I've met friends "
            "through friends, I've done online dating and I just feel like most of the "
            "guys that I meet aren't serious about a relationship."
        ),
        segments=[
            TranscriptionSegment(
                start_seconds=0.0,
                end_seconds=8.0,
                text=(
                    "So I've been single for about four years now and I find it hard "
                    "to meet a guy especially in"
                ),
            ),
            TranscriptionSegment(
                start_seconds=8.0,
                end_seconds=9.5,
                text="California.",
                words=[WordTimestamp(8.0, 9.5, "California.")],
            ),
            TranscriptionSegment(
                start_seconds=9.5,
                end_seconds=20.0,
                text=(
                    "I've tried dating apps, I've met friends through friends, "
                    "I've done online dating"
                ),
            ),
            TranscriptionSegment(
                start_seconds=20.0,
                end_seconds=24.0,
                text=(
                    "and I just feel like most of the guys that I meet aren't serious "
                    "about a relationship."
                ),
            ),
        ],
    )

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

    assert document.algorithm_version == MOSS_ALGORITHM_VERSION
    assert [sentence.text for sentence in document.sentences] == [
        (
            "So I've been single for about four years now and I find it hard to meet "
            "a guy especially in California."
        ),
        (
            "I've tried dating apps, I've met friends through friends, I've done "
            "online dating and I just feel like most of the guys that I meet aren't "
            "serious about a relationship."
        ),
    ]
    assert document.sentences[0].start_ms == 0
    assert document.sentences[0].end_ms == 9500
    assert document.sentences[1].start_ms == 9500
    assert document.sentences[1].end_ms == 24000
    assert all(sentence.reference_speech_duration_ms > 0 for sentence in document.sentences)


def test_split_sentence_text_uses_whisper_word_timestamps(tmp_path):
    transcript = Transcript(
        text="Alpha beta. Gamma delta.",
        segments=[
            TranscriptionSegment(
                start_seconds=0.0,
                end_seconds=1.0,
                text="Alpha beta. Gamma delta.",
                words=[
                    WordTimestamp(0.0, 0.25, "Alpha"),
                    WordTimestamp(0.25, 0.5, "beta."),
                    WordTimestamp(0.5, 0.75, "Gamma"),
                    WordTimestamp(0.75, 1.0, "delta."),
                ],
            )
        ],
    )
    audio_path = tmp_path / "sentence.wav"
    audio_path.write_bytes(_wav_bytes(_speech(1.0), SAMPLE_RATE))
    source = SentenceBoundary(
        index=2,
        start_ms=1000,
        end_ms=5000,
        text="Alpha beta Gamma delta",
        language="en",
        reference_speech_duration_ms=3800,
    )

    replacements = split_sentence_text(
        transcriber=WordTimestampTranscriber(transcript),
        media_path=audio_path,
        sentence=source,
        text="Alpha beta\nGamma delta",
        language="en",
    )

    assert [(item.text, item.start_ms, item.end_ms) for item in replacements] == [
        ("Alpha beta", 1000, 3000),
        ("Gamma delta", 3000, 5000),
    ]