"""Tests for the SGLang-Omni TTS client (converter/clients/sglomni.py)."""

import base64
import io
import json
import tempfile
import unittest
import urllib.error
import wave
from pathlib import Path
from unittest.mock import MagicMock, patch

from converter import config
from converter.clients import SgOmniTTSClient
from converter.clients.base import NonRetryableTTSError
from converter.clients.sglomni import _data_url, _is_loopback
from converter.clients.speakers import QWEN3_TTS_SPEAKERS


class CatalogConsistencyTests(unittest.TestCase):
    """The backend catalog's vendored facts match the converter's."""

    def test_customvoice_speakers_match_the_qwen_table(self):
        from backends.sglomni.catalog import QWEN_CUSTOMVOICE_SPEAKERS
        self.assertEqual(QWEN_CUSTOMVOICE_SPEAKERS, QWEN3_TTS_SPEAKERS)

_DUMMY_CHUNKS = Path(tempfile.gettempdir()) / "audiobook_sglomni_test_chunks"


def _make_wav() -> bytes:
    """A real minimal RIFF/WAVE file (what a server response looks like)."""
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(24000)
        wav_file.writeframes(b"\x01\x00" * 16)
    return buffer.getvalue()


_WAV_BYTES = _make_wav()
_WAV_FRAMES = b"\x01\x00" * 16


class LoopbackTests(unittest.TestCase):
    def test_loopback_hosts(self):
        self.assertTrue(_is_loopback("http://127.0.0.1:8100"))
        self.assertTrue(_is_loopback("http://localhost:8100"))
        self.assertFalse(_is_loopback("http://10.20.30.40:8100"))

    def test_data_url_carries_mime_and_bytes(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "ref.wav"
            path.write_bytes(b"abc")
            url = _data_url(path)
        self.assertTrue(url.startswith("data:audio/wav;base64,"))
        self.assertEqual(
            base64.b64decode(url.partition(";base64,")[2]), b"abc")


class ConnectInputTests(unittest.TestCase):
    """Capability-driven validation before any HTTP is attempted."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.ref = Path(self._tmp.name) / "narrator.wav"
        self.ref.write_bytes(b"abc")
        self.addCleanup(self._tmp.cleanup)

    def _client(self, model="higgs_audio_v3_tts", **kwargs):
        # Bypass _connect (HTTP) — these tests cover the input checks.
        with patch.object(SgOmniTTSClient, "_connect"):
            return SgOmniTTSClient(_DUMMY_CHUNKS, model=model, **kwargs)

    def test_unknown_model_raises(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="nope")
        self.assertIn("Unknown SGLang-Omni model", str(ctx.exception))

    def test_design_model_requires_instructions(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="qwen3_tts_1_7b_voicedesign")
        self.assertIn("--instructions", str(ctx.exception))

    def test_reference_required_model_refuses_to_connect_without_one(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="qwen3_tts_1_7b_base")
        self.assertIn("requires reference audio", str(ctx.exception))

    def test_missing_reference_file_raises(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="higgs_audio_v3_tts",
                         ref_audio=str(Path(self._tmp.name) / "gone.wav"))
        self.assertIn("Reference audio not found", str(ctx.exception))

    def test_clone_capable_model_allows_text_only(self):
        client = self._client(model="higgs_audio_v3_tts")
        self.assertIsNone(client.ref_audio)

    def test_speaker_model_ignores_the_clone_reference(self):
        client = self._client(model="qwen3_tts_0_6b_customvoice",
                              ref_audio=str(self.ref))
        self.assertIsNone(client.ref_audio)

    def test_seed_only_sent_for_models_that_accept_it(self):
        with patch("converter.clients.sglomni.resolve_request_seed",
                   return_value=42):
            client = self._client(model="qwen3_tts_1_7b_base",
                                  ref_audio=str(self.ref))
            self.assertEqual(client._seed, 42)
            client = self._client(model="higgs_audio_v3_tts")
            self.assertIsNone(client._seed)

    def test_negative_seed_is_not_sent(self):
        with patch("converter.clients.sglomni.resolve_request_seed",
                   return_value=-1):
            client = self._client(model="qwen3_tts_1_7b_base",
                                  ref_audio=str(self.ref))
            self.assertIsNone(client._seed)


class ConnectHealthTests(unittest.TestCase):
    """_connect gates on /health and the hosted model."""

    def _response(self, payload):
        response = MagicMock()
        response.__enter__.return_value = response
        response.read.return_value = json.dumps(payload).encode("utf-8")
        return response

    def _connect(self, payloads, **kwargs):
        # urlopen is called once per _get_json call, in order.
        responses = [self._response(payload) for payload in payloads]
        with patch("converter.clients.sglomni.urllib.request.urlopen",
                   side_effect=responses):
            with patch.object(SgOmniTTSClient, "_resolve_reference_text"):
                return SgOmniTTSClient(_DUMMY_CHUNKS,
                                       model="higgs_audio_v3_tts", **kwargs)

    def test_unreachable_server_raises_with_guidance(self):
        import urllib.error
        with patch("converter.clients.sglomni.urllib.request.urlopen",
                   side_effect=urllib.error.URLError("refused")):
            with self.assertRaises(RuntimeError) as ctx:
                SgOmniTTSClient(_DUMMY_CHUNKS, model="higgs_audio_v3_tts")
        self.assertIn("not reachable", str(ctx.exception))
        self.assertIn("sgl-omni", str(ctx.exception))

    def test_booting_server_raises(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._connect([{"status": "unhealthy"}])
        self.assertIn("not healthy", str(ctx.exception))

    def test_foreign_hosted_model_raises_with_both_names(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._connect([
                {"status": "healthy", "stages": []},
                {"data": [{"id": "Zyphra/zonos2"}]},
            ])
        message = str(ctx.exception)
        self.assertIn("Zyphra/zonos2", message)
        self.assertIn("bosonai/higgs-audio-v3-tts-4b", message)

    def test_matching_model_connects(self):
        client = self._connect([
            {"status": "healthy", "stages": []},
            {"data": [{"id": "bosonai/higgs-audio-v3-tts-4b"}]},
        ])
        self.assertEqual(client.entry.repo, "bosonai/higgs-audio-v3-tts-4b")


class PayloadTests(unittest.TestCase):
    """The /v1/audio/speech request shape per voice capability."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.ref = Path(self._tmp.name) / "narrator.wav"
        self.ref.write_bytes(b"abc")
        self.addCleanup(self._tmp.cleanup)

    def _make_client(self, model, **kwargs):
        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
        from backends.sglomni.catalog import entry_by_key
        client.entry = entry_by_key(model)
        client.api_url = "http://127.0.0.1:8100"
        client.voice = kwargs.get("voice")
        if "ref_audio" in kwargs:
            kwargs["ref_audio"] = str(self.ref)
        client.ref_audio = kwargs.get("ref_audio")
        client.ref_text = kwargs.get("ref_text", "")
        client.instructions = kwargs.get("instructions", "")
        client.language = "English"
        client._seed = None
        return client

    def test_speaker_payload_sends_the_preset_name(self):
        client = self._make_client("qwen3_tts_0_6b_customvoice",
                                   voice="Vivian")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["voice"], "Vivian")
        self.assertEqual(payload["model"],
                         "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
        self.assertEqual(payload["response_format"], "wav")
        self.assertNotIn("ref_audio", payload)
        self.assertNotIn("task_type", payload)

    def test_speaker_without_voice_uses_the_server_default(self):
        client = self._make_client("voxtral_tts")
        self.assertEqual(client._request_payload("Hello.")["voice"],
                         "default")

    def test_design_payload_sends_task_type_and_instructions(self):
        client = self._make_client("qwen3_tts_1_7b_voicedesign",
                                   instructions="A warm narrator.")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["task_type"], "VoiceDesign")
        self.assertEqual(payload["instructions"], "A warm narrator.")

    def test_clone_payload_sends_reference_path_on_loopback(self):
        client = self._make_client("higgs_audio_v3_tts",
                                   ref_audio=str(self.ref),
                                   ref_text="A transcript.")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["ref_audio"], str(self.ref.resolve()))
        self.assertEqual(payload["ref_text"], "A transcript.")

    def test_clone_payload_inlines_audio_for_remote_servers(self):
        client = self._make_client("higgs_audio_v3_tts",
                                   ref_audio=str(self.ref))
        client.api_url = "http://10.20.30.40:8100"
        payload = client._request_payload("Hello.")
        self.assertTrue(payload["ref_audio"].startswith(
            "data:audio/wav;base64,"))
        self.assertEqual(
            base64.b64decode(payload["ref_audio"].partition(";base64,")[2]),
            b"abc")
        self.assertNotIn("ref_text", payload)

    def test_clone_without_reference_sends_no_reference_fields(self):
        client = self._make_client("higgs_audio_v3_tts")
        payload = client._request_payload("Hello.")
        self.assertNotIn("ref_audio", payload)
        self.assertEqual(payload["voice"], "default")

    def test_seed_included_when_resolved(self):
        client = self._make_client("qwen3_tts_1_7b_base",
                                   ref_audio="x.wav")
        client._seed = 7
        self.assertEqual(client._request_payload("Hello.")["seed"], 7)

    def test_zonos2_payload_raises_the_generation_cap(self):
        """Zonos2's 1024-frame engine default caps a request at ~12 s."""
        client = self._make_client("zonos2")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["max_new_tokens"], 12288)

    def test_higgs_payload_raises_the_generation_cap(self):
        """Higgs's 2048-frame engine default caps a request at ~27 s
        (75 fps), below a full 250-word sub-chunk."""
        client = self._make_client("higgs_audio_v3_tts")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["max_new_tokens"], 12288)

    def test_models_without_a_cap_send_no_max_new_tokens(self):
        client = self._make_client("moss_tts")
        self.assertNotIn("max_new_tokens",
                         client._request_payload("Hello."))


class RequestErrorTests(unittest.TestCase):
    """OpenAI-style error envelopes decide retryability."""

    def _make_client(self):
        return SgOmniTTSClient.__new__(SgOmniTTSClient)

    def test_bad_request_envelope_is_not_retryable(self):
        client = self._make_client()
        detail = json.dumps({"error": {
            "message": "voice 'nope' not found",
            "type": "BadRequestError", "code": 400}})
        error = client._request_error(400, detail)
        self.assertIsInstance(error, NonRetryableTTSError)
        self.assertIn("voice 'nope' not found", str(error))

    def test_server_error_is_retryable(self):
        client = self._make_client()
        error = client._request_error(503, "overloaded")
        self.assertNotIsInstance(error, NonRetryableTTSError)

    def test_non_json_4xx_is_not_retryable(self):
        client = self._make_client()
        error = client._request_error(422, "plain text rejection")
        self.assertIsInstance(error, NonRetryableTTSError)


class GenerateChunkTests(unittest.TestCase):
    """Chunk generation: WAV output, sub-chunking, bookkeeping."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self._sleep = patch("converter.clients.base.time.sleep")
        self._sleep.start()
        self.addCleanup(self._sleep.stop)
        self.addCleanup(self._tmp.cleanup)

    def _make_client(self):
        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
        from backends.sglomni.catalog import entry_by_key
        client.entry = entry_by_key("higgs_audio_v3_tts")
        client.chunks_dir = Path(self._tmp.name)
        client.api_url = "http://127.0.0.1:8100"
        client.voice = None
        client.ref_audio = None
        client.ref_text = ""
        client.instructions = ""
        client.language = "English"
        client._seed = None
        return client

    def _read_wav(self, path):
        with wave.open(str(path), "rb") as wav_file:
            return wav_file.readframes(wav_file.getnframes())

    def test_generate_chunk_writes_the_wav_response(self):
        client = self._make_client()
        with patch.object(client, "_request_wav", return_value=_WAV_BYTES):
            result = client.generate_chunk("Hello world.", 1)
        self.assertIsNotNone(result)
        path = Path(result)
        self.assertEqual(path.name, "chunk_0001.wav")
        self.assertEqual(self._read_wav(path), _WAV_FRAMES)

    def test_long_text_is_subchunked_and_concatenated(self):
        client = self._make_client()
        text = " ".join(f"word{i}" for i in range(24))
        responses = [_WAV_BYTES, _WAV_BYTES, _WAV_BYTES]
        with patch.object(config, "CHUNK_SIZE", 10), \
                patch.object(client, "_request_wav",
                             side_effect=responses) as mock_wav, \
                patch("converter.clients.sglomni.concat_audio_files") as mock_concat:
            result = client.generate_chunk(text, 1)
        # 24 words at CHUNK_SIZE 10 -> three sub-requests (10/10/4).
        self.assertEqual(mock_wav.call_count, 3)
        self.assertIsNotNone(result)
        mock_concat.assert_called_once()
        args = mock_concat.call_args[0]
        self.assertEqual(len(args[0]), 3)
        self.assertEqual(args[1], Path(result))

    def test_single_subchunk_skips_concatenation(self):
        client = self._make_client()
        with patch.object(client, "_request_wav", return_value=_WAV_BYTES), \
                patch("converter.clients.sglomni.concat_audio_files") as mock_concat:
            client.generate_chunk("Hello.", 1)
        mock_concat.assert_not_called()

    def test_empty_text_fails_the_chunk(self):
        client = self._make_client()
        with patch.object(client, "_request_wav") as mock_wav:
            self.assertIsNone(client.generate_chunk("   ", 1))
        mock_wav.assert_not_called()

    def test_request_failure_fails_the_chunk_attempt(self):
        client = self._make_client()
        with patch.object(client, "_request_wav",
                          side_effect=RuntimeError("down")) as mock_wav:
            self.assertIsNone(client.generate_chunk("Hello.", 1))
        self.assertEqual(mock_wav.call_count, 1)

    def test_stale_chunk_files_are_removed(self):
        stale = Path(self._tmp.name) / "chunk_0001.mp3"
        stale.write_bytes(b"old")
        client = self._make_client()
        with patch.object(client, "_request_wav", return_value=_WAV_BYTES):
            client.generate_chunk("Hello.", 1)
        remaining = sorted(path.name for path in
                           Path(self._tmp.name).glob("chunk_0001.*"))
        self.assertEqual(remaining, ["chunk_0001.wav"])


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