import io
import wave

import numpy as np
import pytest

from sentence_api.transcription import (
    MossTranscriber,
    TranscriptionSegment,
    WordTimestamp,
    _parse_json_segments,
    _plan_chunks,
    _parse_word_timestamps,
    split_sentences_at_punctuation,
    _wav_bytes,
    _wav_duration_seconds,
)


SAMPLE_RATE = 16000


def _speech(seconds: float, sample_rate: int = SAMPLE_RATE) -> 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, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    return np.zeros(int(seconds * sample_rate))


def _audio(*parts: np.ndarray, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    return np.concatenate([np.asarray(part) for part in parts]).astype(np.float32)


def test_plan_chunks_cuts_at_silence_gaps():
    samples = _audio(
        _speech(10),
        _silence(0.8),
        _speech(29.2),
        _silence(1.0),
        _speech(29.0),
        _silence(0.6),
        _speech(24.4),
    )
    chunks = _plan_chunks(samples, SAMPLE_RATE)
    assert chunks[0][0] == pytest.approx(0.0)
    assert chunks[-1][1] == pytest.approx(95.0)
    for (start, end), (next_start, _) in zip(chunks, chunks[1:]):
        assert end == pytest.approx(next_start)
        assert end - start <= 34.0
    assert 10.0 <= chunks[0][1] <= 10.8
    assert 40.0 <= chunks[1][1] <= 41.0
    assert 70.0 <= chunks[2][1] <= 70.6


def test_plan_chunks_falls_back_to_fixed_windows():
    chunks = _plan_chunks(_speech(95), SAMPLE_RATE)
    assert chunks == [(0.0, 30.0), (30.0, 60.0), (60.0, 90.0), (90.0, 95.0)]


def test_plan_chunks_merges_short_tail():
    chunks = _plan_chunks(_speech(93), SAMPLE_RATE)
    assert chunks == [(0.0, 30.0), (30.0, 60.0), (60.0, 93.0)]


def test_parse_json_segments_filters_hallucinated_segments():
    raw = [
        {"start": 0.0, "end": 3.0, "text": "hello world", "compression_ratio": 1.4},
        {"start": 3.5, "end": 6.0, "text": "hello world", "compression_ratio": 12.0},
        {"start": 6.5, "end": 9.0, "text": "no ratio field"},
    ]
    filtered = _parse_json_segments(raw, compression_limit=2.4)
    assert [segment.text for segment in filtered] == ["hello world", "no ratio field"]
    assert len(_parse_json_segments(raw)) == 3


def test_wav_bytes_roundtrip(tmp_path):
    samples = _speech(2.0)
    data = _wav_bytes(samples, SAMPLE_RATE)
    assert data.startswith(b"RIFF")
    path = tmp_path / "chunk.wav"
    path.write_bytes(data)
    assert _wav_duration_seconds(path) == pytest.approx(2.0, abs=0.05)
    with wave.open(str(path), "rb") as wav:
        assert wav.getnchannels() == 1
        assert wav.getsampwidth() == 2
        assert wav.getframerate() == SAMPLE_RATE


def test_wav_duration_returns_none_for_non_wav(tmp_path):
    path = tmp_path / "note.txt"
    path.write_text("not audio")
    assert _wav_duration_seconds(path) is None


def test_chunked_transcribe_offsets_timestamps(monkeypatch):
    requests = []
    call_count = {"n": 0}

    def fake_post(endpoint, data=None, files=None, timeout=None):
        call_count["n"] += 1
        requests.append(data)
        audio = files["file"][1].read()
        with wave.open(io.BytesIO(audio), "rb") as wav:
            duration = wav.getnframes() / wav.getframerate()
        if call_count["n"] == 1:
            segments = [
                {"start": 1.0, "end": 4.0, "text": "first part", "compression_ratio": 1.4},
                {"start": 27.0, "end": 30.0, "text": "second part", "compression_ratio": 1.5},
            ]
        elif call_count["n"] == 2:
            segments = [
                {
                    "start": 1.0,
                    "end": 4.0,
                    "text": "third part",
                    "compression_ratio": 1.4,
                    "words": [
                        {"word": "third", "start": 1.0, "end": 2.5},
                        {"word": " part", "start": 2.6, "end": 4.0},
                    ],
                },
                {"start": 27.0, "end": 30.0, "text": "fourth part", "compression_ratio": 1.5},
            ]
        else:
            segments = [{"start": 0.5, "end": 3.5, "text": "tail part", "compression_ratio": 1.3}]
        assert abs(duration - 30.0) < 1.0 or abs(duration - 5.0) < 1.0

        class Response:
            status_code = 200

            def raise_for_status(self):
                pass

            def json(self):
                return {"text": "", "segments": segments}

        return Response()

    monkeypatch.setattr("sentence_api.transcription.httpx.post", fake_post)
    transcriber = MossTranscriber(endpoint="http://whisper:9000", model="whisper")
    result = transcriber._transcribe_chunked(_speech(65), SAMPLE_RATE, None)
    assert [(s.start_seconds, s.end_seconds, s.text) for s in result.segments] == [
        (1.0, 4.0, "first part"),
        (27.0, 30.0, "second part"),
        (31.0, 34.0, "third part"),
        (57.0, 60.0, "fourth part"),
        (60.5, 63.5, "tail part"),
    ]
    assert result.text == "first part second part third part fourth part tail part"
    assert [(w.start_seconds, w.end_seconds, w.text) for w in result.segments[2].words] == [
        (31.0, 32.5, "third"),
        (32.6, 34.0, "part"),
    ]
    assert result.segments[0].words is None
    assert len(requests) == 3
    assert all(request.get("condition_on_previous_text") == "false" for request in requests)
    assert all(request.get("timestamp_granularities[]") == "word" for request in requests)


def test_parse_json_segments_parses_word_timestamps():
    raw = [
        {
            "start": 0.0,
            "end": 4.0,
            "text": "Hello world. Good day.",
            "words": [
                {"word": "Hello", "start": 0.0, "end": 0.6},
                {"word": " world.", "start": 0.7, "end": 1.5},
                {"word": " Good", "start": 1.8, "end": 2.4},
                {"word": " day.", "start": 2.5, "end": 3.2},
            ],
        },
        {"start": 4.5, "end": 6.0, "text": "no words"},
    ]
    segments = _parse_json_segments(raw)
    assert [word.text for word in segments[0].words] == ["Hello", "world.", "Good", "day."]
    assert segments[1].words is None


def test_parse_word_timestamps_ignores_invalid_entries():
    words = _parse_word_timestamps(
        [
            {"word": "ok", "start": 0.0, "end": 0.5},
            {"word": "bad"},
            {"word": "", "start": 1.0, "end": 1.5},
            {"word": "flat", "start": 2.0, "end": 2.0},
        ]
    )
    assert [(word.text, word.start_seconds, word.end_seconds) for word in words] == [
        ("ok", 0.0, 0.5)
    ]
    assert _parse_word_timestamps(None) is None
    assert _parse_word_timestamps("nope") is None


def test_split_sentences_merges_segments_without_punctuation():
    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."
            ),
            words=[WordTimestamp(20.0, 24.0, "relationship.")],
        ),
    ]
    sentences = split_sentences_at_punctuation(segments)
    assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [
        (
            "So I've been single for about four years now and I find it hard to meet "
            "a guy especially in California.",
            0.0,
            9.5,
        ),
        (
            "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.",
            9.5,
            24.0,
        ),
    ]


def test_split_sentences_uses_word_timestamps_within_segment():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=4.0,
        text="Hello world. Good day. Nice to meet you.",
        words=[
            WordTimestamp(0.0, 0.6, "Hello"),
            WordTimestamp(0.7, 1.5, "world."),
            WordTimestamp(1.8, 2.4, "Good"),
            WordTimestamp(2.5, 3.2, "day."),
            WordTimestamp(3.3, 3.8, "Nice"),
            WordTimestamp(3.9, 4.0, "you."),
        ],
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [
        ("Hello world.", 0.0, 1.5),
        ("Good day.", 1.5, 3.2),
        ("Nice to meet you.", 3.2, 4.0),
    ]


def test_split_sentences_splits_at_question_marks():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=10.0,
        text="Really? Wait a minute. What about now?",
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [s.text for s in sentences] == [
        "Really?",
        "Wait a minute.",
        "What about now?",
    ]
    assert sentences[0].start_seconds == 0.0
    assert sentences[0].end_seconds == sentences[1].start_seconds
    assert sentences[1].end_seconds == sentences[2].start_seconds
    assert sentences[2].end_seconds == 10.0


def test_split_sentences_splits_at_comma_after_fifty_words():
    before_comma = " ".join(f"word{index}" for index in range(1, 51))
    after_comma = "This continuation becomes its own sentence."
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=55.0,
        text=f"{before_comma}, {after_comma}",
    )

    sentences = split_sentences_at_punctuation([segment])

    assert [s.text for s in sentences] == [
        f"{before_comma},",
        after_comma,
    ]
    assert sentences[0].start_seconds == 0.0
    assert sentences[0].end_seconds == sentences[1].start_seconds
    assert sentences[1].end_seconds == 55.0


def test_split_sentences_merges_comma_piece_shorter_than_fifty_words():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=5.0,
        text="One, two, and three.",
    )

    sentences = split_sentences_at_punctuation([segment])

    assert [s.text for s in sentences] == ["One, two, and three."]


def test_split_sentences_falls_back_to_proportional():
    segments = [
        TranscriptionSegment(
            start_seconds=10.0, end_seconds=14.0, text="This is the first"
        ),
        TranscriptionSegment(
            start_seconds=14.0,
            end_seconds=18.0,
            text="sentence. And this is the second one.",
        ),
    ]
    sentences = split_sentences_at_punctuation(segments)
    assert [s.text for s in sentences] == [
        "This is the first sentence.",
        "And this is the second one.",
    ]
    assert sentences[0].start_seconds == 10.0
    assert sentences[0].end_seconds == sentences[1].start_seconds
    assert sentences[1].end_seconds == 18.0


def test_split_sentences_falls_back_when_word_ends_are_invalid():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=2.0,
        text="One. Two.",
        words=[
            WordTimestamp(0.0, 2.5, "One."),
            WordTimestamp(2.6, 3.0, "Two."),
        ],
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [s.text for s in sentences] == ["One.", "Two."]
    assert sentences[0].end_seconds == sentences[1].start_seconds
    assert sentences[1].end_seconds == 2.0


def test_split_sentences_keeps_incomplete_tail_as_sentence():
    segment = TranscriptionSegment(
        start_seconds=1.0, end_seconds=2.0, text="no period here"
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [
        ("no period here", 1.0, 2.0)
    ]
    assert split_sentences_at_punctuation([]) == []


def test_split_sentences_handles_ellipsis_and_dots_only():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=2.0,
        text="Wait... What? ...",
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [s.text for s in sentences] == ["Wait...", "What?"]
    assert sentences[0].start_seconds == 0.0
    assert sentences[0].end_seconds == sentences[1].start_seconds
    assert sentences[1].end_seconds == 2.0
    assert split_sentences_at_punctuation(
        [TranscriptionSegment(0.0, 1.0, "...")]
    ) == []


def test_split_sentences_ignores_decimal_points():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=6.0,
        text="The rate rose by 0.4%. That is a big jump.",
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [s.text for s in sentences] == [
        "The rate rose by 0.4%.",
        "That is a big jump.",
    ]


def test_split_sentences_ignores_am_pm_abbreviations():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=8.0,
        text=(
            "It was at 9 o'clock p.m. Eastern, President Trump was speaking. "
            "The room was quiet."
        ),
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [s.text for s in sentences] == [
        "It was at 9 o'clock p.m. Eastern, President Trump was speaking.",
        "The room was quiet.",
    ]

    morning = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=4.0,
        text="I woke up at 7 a.m. and had coffee. Then I left.",
    )
    sentences = split_sentences_at_punctuation([morning])
    assert [s.text for s in sentences] == [
        "I woke up at 7 a.m. and had coffee.",
        "Then I left.",
    ]


def test_split_sentences_merges_segment_ending_with_abbreviation():
    segments = [
        TranscriptionSegment(0.0, 5.0, "The meeting ends at 5 p.m."),
        TranscriptionSegment(5.0, 9.0, "Eastern time. Then we go home."),
    ]
    sentences = split_sentences_at_punctuation(segments)
    assert [s.text for s in sentences] == [
        "The meeting ends at 5 p.m. Eastern time.",
        "Then we go home.",
    ]


def test_split_sentences_distinguishes_no_reply_from_number():
    reply = TranscriptionSegment(0.0, 4.0, "No. I don't think so.")
    assert [s.text for s in split_sentences_at_punctuation([reply])] == [
        "No.",
        "I don't think so.",
    ]
    numbered = TranscriptionSegment(0.0, 4.0, "See no. 5 on the list.")
    assert [s.text for s in split_sentences_at_punctuation([numbered])] == [
        "See no. 5 on the list."
    ]


def test_split_sentences_ignores_common_abbreviations():
    segment = TranscriptionSegment(
        start_seconds=0.0,
        end_seconds=4.0,
        text="Dr. Smith said the U.S. economy is growing. He was right.",
    )
    sentences = split_sentences_at_punctuation([segment])
    assert [s.text for s in sentences] == [
        "Dr. Smith said the U.S. economy is growing.",
        "He was right.",
    ]


def test_post_audio_requests_word_timestamps_and_falls_back(monkeypatch):
    calls = []

    class RejectedResponse:
        status_code = 400

        def raise_for_status(self):
            pass

        def json(self):
            return {"error": "unknown parameter"}

    class OkResponse:
        status_code = 200

        def raise_for_status(self):
            pass

        def json(self):
            return {"text": "hi.", "segments": []}

    def fake_post(endpoint, data=None, files=None, timeout=None):
        calls.append(dict(data))
        if "timestamp_granularities[]" in data:
            return RejectedResponse()
        return OkResponse()

    monkeypatch.setattr("sentence_api.transcription.httpx.post", fake_post)
    transcriber = MossTranscriber(endpoint="http://whisper:9000", model="whisper")
    result = transcriber._post_audio(
        io.BytesIO(b"fake-audio"), "clip.wav", "audio/wav", "en"
    )
    assert result["text"] == "hi."
    assert calls[0]["timestamp_granularities[]"] == "word"
    assert "timestamp_granularities[]" not in calls[1]
