"""Regression cases for silent, stalled, and unterminated model streams."""
import asyncio
import base64
from contextlib import aclosing
import json

import numpy as np
import pytest

from test_gradio_ui import anyio_backend, studio, browser_request, body, EngineClient
import server
from accounts import Identity
from gradio_ui import events


def pcm(seconds, amplitude=0, rate=44100):
    return np.full(round(seconds * rate), amplitude, dtype="<i2").tobytes()


def delta(audio):
    return "data: " + json.dumps({"type": "speech.audio.delta", "audio": base64.b64encode(audio).decode()})


class Response:
    status_code = 200
    headers = {"x-sample-rate": "44100"}

    def __init__(self, chunks, finish="stop", hold=False, heartbeat=False):
        self.chunks, self.finish, self.hold, self.heartbeat = chunks, finish, hold, heartbeat
        self.closed = self.cancelled = False

    async def aiter_lines(self):
        try:
            for chunk in self.chunks:
                yield delta(chunk)
            if self.finish is not None:
                yield 'data: ' + json.dumps({"type": "speech.audio.done", "finish_reason": self.finish})
            if self.heartbeat:
                while True:
                    yield ': heartbeat'
                    await asyncio.sleep(.001)
            if self.hold:
                await asyncio.Event().wait()
        finally:
            self.cancelled = True

    async def aclose(self):
        self.closed = True


class Client(EngineClient):
    def __init__(self, *responses):
        self.responses = list(responses)
        self.requests = []

    async def send(self, request, stream):
        self.requests.append(self.payload)
        return self.responses[len(self.requests) - 1]


async def generate(studio):
    response = await server.generate(body(studio), Identity(keys=("test",), name="test", signed_in=False))
    return [e async for e in events(response)]


def assert_released(*responses):
    assert all(r.closed and r.cancelled for r in responses)
    assert server.gate.active == 0 and not server.limits.active


def test_quiet_edges_and_internal_pauses_are_trimmed_without_losing_soft_speech():
    rate = 1000
    guard = server.SilenceGuard(rate)
    # Soft speech just above the existing threshold must survive unchanged.
    speech = pcm(.5, 150, rate)
    raw = pcm(1, rate=rate) + speech + pcm(2, rate=rate) + speech + pcm(1, rate=rate)
    output = guard.feed(raw) + guard.feed(b"", final=True)
    assert output == speech + pcm(.75, rate=rate) + speech + pcm(.15, rate=rate)
    assert guard.speech_detected and not guard.stalled
    assert guard.dropped == 3100


def test_stall_is_sticky_even_when_speech_follows_in_same_chunk():
    guard = server.SilenceGuard(1000)
    speech = pcm(.5, 3000, 1000)
    output = guard.feed(speech + pcm(4.5, rate=1000) + speech)
    output += guard.feed(b"", final=True)
    assert guard.stalled and output == speech + pcm(.15, rate=1000)
    assert not guard.feed(speech)


def test_all_silence_never_produces_audio():
    guard = server.SilenceGuard(1000)
    assert not guard.feed(pcm(4, rate=1000))
    assert not guard.feed(b"", final=True)
    assert guard.stalled and not guard.speech_detected
    assert guard.dropped == 4000


def test_arbitrary_byte_boundaries_and_partial_final_frame_preserve_speech():
    audio = pcm(.123, 1000, 1000)
    guard = server.SilenceGuard(1000)
    chunks = [audio[:1], audio[1:53], audio[53:177], audio[177:]]
    output = b"".join(guard.feed(chunk) for chunk in chunks) + guard.feed(b"", final=True)
    assert output == audio


@pytest.mark.anyio
async def test_silent_take_retries_once_before_playing(studio, monkeypatch):
    silent, spoken = Response([pcm(4)]), Response([pcm(.5, 3000)])
    client = Client(silent, spoken)
    monkeypatch.setattr(server, "http", client)
    result = await generate(studio)
    assert [e['type'] for e in result] == ['retrying', 'start', 'audio', 'done']
    assert len(client.requests) == 2
    assert client.requests[0] == client.requests[1]
    assert result[-1]['retry_count'] == 1 and result[-1]['audio_s'] == .5
    assert len(server.store.takes) == 1
    assert_released(silent, spoken)


@pytest.mark.anyio
async def test_two_silent_attempts_fail_without_saved_take(studio, monkeypatch):
    first, second = Response([pcm(.5)]), Response([pcm(4)])
    monkeypatch.setattr(server, "http", Client(first, second))
    result = await generate(studio)
    assert [e['type'] for e in result] == ['retrying', 'error']
    assert 'two attempts' in result[-1]['message']
    assert not server.store.takes
    assert_released(first, second)


@pytest.mark.anyio
async def test_partial_speech_stall_closes_request_and_warns_without_retry(studio, monkeypatch):
    response = Response([pcm(3.5, 3000), pcm(6)], finish=None, hold=True)
    client = Client(response)
    monkeypatch.setattr(server, "http", client)
    result = await asyncio.wait_for(generate(studio), 1)
    assert result[-1]['finish_reason'] == 'stall' and result[-1]['warning']
    assert result[-1]['audio_s'] == 3.65 and result[-1]['retry_count'] == 0
    assert len(client.requests) == 1
    assert_released(response)


@pytest.mark.anyio
async def test_done_does_not_wait_for_connection_eof(studio, monkeypatch):
    response = Response([pcm(.5, 3000)], hold=True)
    monkeypatch.setattr(server, "http", Client(response))
    result = await asyncio.wait_for(generate(studio), 1)
    assert result[-1]['finish_reason'] == 'stop' and not result[-1]['warning']
    assert_released(response)


@pytest.mark.anyio
async def test_large_opening_streams_small_packets_without_changing_audio_or_retry(studio, monkeypatch):
    # Preserve the validated opening, but a player need not receive its entire
    # multi-second base64 line before it can decode the first audio packet.
    raw = pcm(3.5, 3000)
    response = Response([raw])
    client = Client(response)
    monkeypatch.setattr(server, 'http', client)
    result = await generate(studio)
    packets = [base64.b64decode(e['pcm']) for e in result if e['type'] == 'audio']
    assert b''.join(packets) == raw
    assert max(map(len, packets)) <= 44100  # 0.5 seconds of mono PCM16.
    assert result[-1]['retry_count'] == 0 and len(client.requests) == 1
    assert sum('ttfa_ms' in e for e in result if e['type'] == 'audio') == 1
    assert_released(response)


@pytest.mark.anyio
@pytest.mark.parametrize('heartbeat', [False, True])
async def test_missing_audio_times_out_even_if_heartbeats_continue(studio, monkeypatch, heartbeat):
    response = Response([pcm(.5, 3000)], finish=None, hold=True, heartbeat=heartbeat)
    monkeypatch.setattr(server, "http", Client(response))
    monkeypatch.setattr(server, "ENGINE_IDLE_TIMEOUT_S", .03)
    result = await asyncio.wait_for(generate(studio), 1)
    assert result[-1]['finish_reason'] == 'timeout' and result[-1]['warning']
    assert result[-1]['audio_s'] == .5 and result[-1]['retry_count'] == 0
    assert_released(response)


@pytest.mark.anyio
async def test_timeout_before_speech_is_free_and_closes_stream(studio, monkeypatch):
    response = Response([], finish=None, hold=True)
    monkeypatch.setattr(server, "http", Client(response))
    monkeypatch.setattr(server, "ENGINE_START_TIMEOUT_S", .03)
    result = await asyncio.wait_for(generate(studio), 1)
    assert result[-1]['type'] == 'error' and 'time' in result[-1]['message']
    assert not server.store.takes
    assert_released(response)


@pytest.mark.anyio
async def test_header_timeout_cancels_engine_and_releases_admission(studio, monkeypatch):
    class WaitingClient(EngineClient):
        cancelled = False
        async def send(self, request, stream):
            try:
                await asyncio.Event().wait()
            finally:
                self.cancelled = True
    client = WaitingClient(None)
    monkeypatch.setattr(server, "http", client)
    monkeypatch.setattr(server, "ENGINE_START_TIMEOUT_S", .03)
    result = await asyncio.wait_for(generate(studio), 1)
    assert result[-1]['type'] == 'error' and client.cancelled
    assert server.gate.active == 0 and not server.limits.active


@pytest.mark.anyio
@pytest.mark.parametrize('finish', ['length', None])
async def test_truncated_take_is_visibly_marked_in_gradio(studio, browser_request, monkeypatch, finish):
    response = Response([pcm(3.5, 3000)], finish=finish)
    monkeypatch.setattr(server, "http", Client(response))
    result = [e async for e in studio.run(body(studio), None, [], browser_request)]
    assert 'incomplete' in result[-1][3]
    assert_released(response)


@pytest.mark.anyio
async def test_tiny_audible_final_frame_still_starts(studio, monkeypatch):
    response = Response([pcm(.025, 3000)])
    monkeypatch.setattr(server, "http", Client(response))
    result = await generate(studio)
    assert [e['type'] for e in result] == ['start', 'audio', 'done']
    assert_released(response)


def test_default_sampling_has_no_top_k_cap(studio):
    """top-k 30 left 11 of 40 H200 takes silent to the length cap; no cap left none."""
    choice = server.VoiceChoice('random')
    assert server.engine_request(choice, "Hello there.", [], body(studio))['top_k'] == server.ALL_CANDIDATES
    assert server.engine_request(choice, "Hello there.", [], body(studio).model_copy(update={'top_k': 40}))['top_k'] == 40
    assert server.GenerateBody(text="Hi.", voice=server.VoiceBody(mode="random")).top_k == 0


def test_text_budget_keeps_room_for_multilingual_and_slow_scripts(studio):
    choice = server.VoiceChoice('random')
    short = body(studio)
    assert 18 <= server.token_budget(short, choice) * 2048 / 44100 < 30
    slow = short.model_copy(update={'text': '[slowly] Hello there.'})
    assert server.token_budget(slow, choice) > server.token_budget(short, choice)
    for text in ['你好世界。' * 30, 'ภาษาไทย' * 40, 'The longer speech continues. ' * 20]:
        assert server.token_budget(short.model_copy(update={'text': text}), choice) == server.MAX_NEW_TOKENS
    directed = short.model_copy(update={'text': 'Direction: calm.\nLine: First sentence.\nSecond sentence.'})
    assert server.spoken_text(directed.text) == 'First sentence. Second sentence.'
    assert server.token_budget(directed, choice) == server.token_budget(short.model_copy(update={'text': 'First sentence. Second sentence.'}), choice)


@pytest.mark.anyio
async def test_stop_during_idle_read_cancels_pending_read_and_closes_response(studio, monkeypatch):
    response = Response([pcm(3.5, 3000)], finish=None, hold=True)
    monkeypatch.setattr(server, "http", Client(response))
    stream_response = await server.generate(body(studio), Identity(keys=("test",), name="test", signed_in=False))
    stream = events(stream_response)
    async with aclosing(stream):
        assert (await anext(stream))['type'] == 'start'
        received = bytearray()
        while len(received) < len(pcm(3.5, 3000)):
            event = await anext(stream)
            assert event['type'] == 'audio'
            received.extend(base64.b64decode(event['pcm']))
        waiting = asyncio.create_task(anext(stream))
        await asyncio.sleep(.01)
        waiting.cancel()
        with pytest.raises(asyncio.CancelledError):
            await waiting
    assert_released(response)


@pytest.mark.anyio
async def test_wall_clock_limit_applies_even_while_audio_keeps_arriving(studio, monkeypatch):
    class EndlessResponse(Response):
        async def aiter_lines(self):
            try:
                while True:
                    yield delta(pcm(.05, 3000))
                    await asyncio.sleep(.005)
            finally:
                self.cancelled = True
    response = EndlessResponse([])
    monkeypatch.setattr(server, "http", Client(response))
    monkeypatch.setattr(server, "TAKE_TIMEOUT_S", .04)
    result = await asyncio.wait_for(generate(studio), 1)
    assert result[-1]['finish_reason'] == 'timeout' and result[-1]['warning']
    assert_released(response)


@pytest.mark.anyio
@pytest.mark.parametrize('fragment_s', [1.3, 2.0, 2.4])
async def test_early_stalled_fragment_is_retried_before_playback(studio, monkeypatch, fragment_s):
    failed = Response([pcm(fragment_s, 3000), pcm(5)], finish=None, hold=True)
    good = Response([pcm(3.5, 1000)])
    client = Client(failed, good)
    monkeypatch.setattr(server, "http", client)
    result = await generate(studio)
    assert [e['type'] for e in result if e['type'] != 'audio'] == ['retrying', 'start', 'done']
    output = b''.join(base64.b64decode(e['pcm']) for e in result if e['type'] == 'audio')
    assert output == pcm(3.5, 1000)
    assert result[-1]['retry_count'] == 1
    assert_released(failed, good)


@pytest.mark.anyio
async def test_opening_plays_at_first_64_frame_window_without_waiting_for_next_chunk(studio, monkeypatch):
    response = Response([pcm(frames * 2048 / 44100, 3000) for frames in [10, 10, 12, 32]],
                        finish=None, hold=True)
    monkeypatch.setattr(server, 'http', Client(response))
    backend = await server.generate(body(studio), Identity(keys=("test",), name="test", signed_in=False))
    async with aclosing(events(backend)) as stream:
        assert (await asyncio.wait_for(anext(stream), .5))['type'] == 'start'
        assert (await anext(stream))['type'] == 'audio'
    assert_released(response)


@pytest.mark.anyio
async def test_two_failed_openings_are_not_played(studio, monkeypatch):
    first = Response([pcm(1.3, 3000), pcm(5)])
    second = Response([pcm(2, 3000)], finish='length')
    monkeypatch.setattr(server, 'http', Client(first, second))
    result = await generate(studio)
    assert [e['type'] for e in result] == ['retrying', 'error']
    assert not server.store.takes
    assert_released(first, second)