import asyncio
import json
import unittest
from unittest.mock import AsyncMock, patch

from pipecat.frames.frames import TTSAudioRawFrame, TTSStoppedFrame
from services.pipecat.dashscope_tts import DashScopeTTSService, SynthesisSession, SUPPORTED_MODELS
from services.model_resource_tester import test_model_resource

URL = 'wss://workspace.cn-beijing.maas.aliyuncs.com/api-ws/v1/inference'


class Socket:
    def __init__(self):
        self.sent = []
        self.queue = asyncio.Queue()
        self.close = AsyncMock()

    async def send(self, raw):
        msg = json.loads(raw)
        self.sent.append(msg)
        if msg['header']['action'] == 'run-task':
            self.task_id = msg['header']['task_id']
            self.event('task-started')

    def event(self, name):
        self.queue.put_nowait(json.dumps({'header': {'event': name, 'task_id': self.task_id}}))

    async def recv(self):
        return await self.queue.get()


class DashScopeTests(unittest.IsolatedAsyncioTestCase):
    async def test_protocol_streams_before_finish_and_preserves_spaces(self):
        socket = Socket()
        session = SynthesisSession(URL, 'secret', 'cosyvoice-v1', 'longxiaochun')
        with patch('services.pipecat.dashscope_tts.connect', AsyncMock(return_value=socket)):
            await session.start()
            await session.text('Hello')
            await session.text(' world')
            socket.queue.put_nowait(b'\x01\x00' * 20)
            audio = session.audio()
            self.assertEqual(await anext(audio), b'\x01\x00' * 20)
            self.assertFalse(session.finished_input)
            await session.finish()
            await session.finish()
            socket.event('task-finished')
            with self.assertRaises(StopAsyncIteration):
                await anext(audio)
            await session.close()
        self.assertEqual([m['header']['action'] for m in socket.sent],
                         ['run-task', 'continue-task', 'continue-task', 'finish-task'])
        self.assertEqual(socket.sent[2]['payload']['input']['text'], ' world')
        self.assertEqual(socket.sent[0]['payload']['parameters']['format'], 'pcm')
        self.assertEqual(len({m['header']['task_id'] for m in socket.sent}), 1)

    async def test_empty_audio_and_provider_failure(self):
        for failure in (False, True):
            socket = Socket()
            session = SynthesisSession(URL, 'secret', 'cosyvoice-v2', 'voice')
            with patch('services.pipecat.dashscope_tts.connect', AsyncMock(return_value=socket)):
                await session.start()
                socket.event('task-failed' if failure else 'task-finished')
                with self.assertRaises(RuntimeError):
                    await anext(session.audio())
                await session.close()

    async def test_pipecat_audio_completion_and_interruption(self):
        socket = Socket()
        service = DashScopeTTSService(api_key='secret', model='qwen-audio-3.1-tts-flash',
                                     voice='voice', url=URL)
        service._sample_rate = 24000
        service.create_task = asyncio.create_task
        service.cancel_task = AsyncMock(side_effect=self.cancel_task)
        service.audio_context_available = lambda _: True
        service.append_to_audio_context = AsyncMock()
        service.remove_audio_context = AsyncMock()
        with patch('services.pipecat.dashscope_tts.connect', AsyncMock(return_value=socket)):
            async for _ in service.run_tts('你好', 'turn1'):
                pass
            receiver = service._streams['turn1'].receiver
            socket.queue.put_nowait(b'\0\0' * 20)
            await service.flush_audio('turn1')
            socket.event('task-finished')
            await receiver
            frames = [call.args[1] for call in service.append_to_audio_context.call_args_list]
            self.assertIsInstance(frames[0], TTSAudioRawFrame)
            self.assertIsInstance(frames[-1], TTSStoppedFrame)
            async for _ in service.run_tts('下一轮', 'turn2'):
                pass
            await service.on_audio_context_interrupted('turn2')
            self.assertFalse(service._streams)
            self.assertGreaterEqual(socket.close.await_count, 2)

    async def cancel_task(self, task):
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

    async def test_tester_requires_audio_and_completion(self):
        with patch.object(SynthesisSession, 'start', AsyncMock()), \
             patch.object(SynthesisSession, 'text', AsyncMock()), \
             patch.object(SynthesisSession, 'finish', AsyncMock()), \
             patch.object(SynthesisSession, 'close', AsyncMock()):
            async def no_audio(self):
                raise RuntimeError('no audio secret')
                yield b''
            with patch.object(SynthesisSession, 'audio', no_audio):
                result = await test_model_resource('dashscope-tts', 'TTS',
                    {'apiUrl': URL, 'modelId': 'cosyvoice-v2', 'voice': 'voice'}, {'apiKey': 'secret'})
            self.assertFalse(result.ok)
            self.assertNotIn('secret', result.detail)

    def test_factory_routes_dashscope_to_token_streaming(self):
        from models import AssistantConfig
        from services.pipecat.service_factory import create_tts
        service = create_tts(AssistantConfig(
            tts_interface_type="dashscope-tts", tts_model="cosyvoice-v1",
            voice="longxiaochun", tts_api_key="secret", tts_base_url=URL))
        self.assertIsInstance(service, DashScopeTTSService)
        self.assertTrue(service._is_streaming_tokens)

    def test_models_and_old_http_url(self):
        for model in SUPPORTED_MODELS:
            SynthesisSession(URL, 'secret', model, 'voice')
        with self.assertRaisesRegex(ValueError, '旧 HTTP'):
            SynthesisSession('https://dashscope.aliyuncs.com/api/v1/services/audio/tts',
                             'secret', 'cosyvoice-v1', 'voice')