"""Regression for stale message intents swallowing the next AI recording."""
import asyncio
import struct
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import Mock

import pytest

from handlers import audio_packet_parser as parser
from handlers import audio_session_handler as audio_handler
from handlers import session_cleanup_handler as cleanup
from handlers import websocket_message_handler as messages
from handlers import websocket_handler as endpoint
from services import device_target_cache as targets_module
from services.audio_session import AudioSessionManager
from services.connection_manager import ConnectionManager
from services.device_target_cache import DeviceTargetCache
from services.target_audio_cache import TargetAudioCache

DEVICE = 'TalkingQ_local_voice'
PROCESS_PARENT_MESSAGE = messages.process_parent_leave_message


def packet(sid, kind, audio=b''):
    return DEVICE.encode()+b'\0'+sid.encode()+b'\0'+struct.pack('<IBI', 1, kind, len(audio))+audio


@pytest.fixture
def state(monkeypatch):
    now = [100.0]
    connections, sessions, buffers = ConnectionManager(), AudioSessionManager(), TargetAudioCache()
    targets = DeviceTargetCache(timeout_seconds=10, clock=lambda: now[0])
    started, ended = [], []
    ws = SimpleNamespace(send_text=AsyncMock(), accept=AsyncMock(), client='local',
                         client_state=SimpleNamespace(name='CONNECTED'))
    for module in [messages, endpoint, targets_module]:
        monkeypatch.setattr(module, 'connection_manager', connections)
    for module in [messages, cleanup]:
        monkeypatch.setattr(module, 'device_target_cache', targets)
    for module in [messages, targets_module]:
        monkeypatch.setattr(module, 'target_audio_cache', buffers)
    for module in [messages, cleanup, parser]:
        monkeypatch.setattr(module, 'audio_session_manager', sessions)
    monkeypatch.setattr(cleanup, 'task_manager', SimpleNamespace(cancel_session_tasks=AsyncMock()))
    monkeypatch.setattr(parser, 'get_asr', AsyncMock(side_effect=lambda _: SimpleNamespace(connected=False)))
    monkeypatch.setattr(parser, 'get_llm', AsyncMock(side_effect=lambda _: SimpleNamespace(close=AsyncMock())))
    monkeypatch.setattr(parser, 'get_tts', AsyncMock(side_effect=lambda _: SimpleNamespace(close=AsyncMock())))

    async def prepare(session, device_id, sid):
        started.append(sid)
        session.asr_service.connected = True
        session.send_queue = asyncio.Queue()

    async def finish(session, key):
        chunks = []
        while not session.send_queue.empty():
            chunks.append(session.send_queue.get_nowait())
        ended.append((key[1], b''.join(chunks)))

    monkeypatch.setattr(audio_handler, 'handle_start_packet', prepare)
    monkeypatch.setattr(audio_handler, 'handle_end_packet', finish)
    monkeypatch.setattr(messages, 'process_parent_leave_message', AsyncMock())
    monkeypatch.setattr(messages, 'process_cached_audio', AsyncMock())
    return SimpleNamespace(now=now, connections=connections, sessions=sessions, buffers=buffers,
                           targets=targets, ws=ws, started=started, ended=ended)


async def send(state, sid, kind, audio=b'', *, ws=None, active=None):
    return await messages.handle_binary_message(ws or state.ws, DEVICE, 'LOCAL-SERIAL',
                                                packet(sid, kind, audio), active, None)


async def set_target(state, kind):
    if kind == 'parent':
        await state.targets.set_parent_target(DEVICE, DEVICE)
    else:
        await state.targets.set_target(DEVICE, 'TalkingQ_receiver')
    return await state.targets.get_voice_target(DEVICE)


@pytest.mark.asyncio
@pytest.mark.parametrize('kind', ['parent', 'device'])
async def test_ai_audio_recovers_from_pending_message_without_losing_first_chunk(state, kind):
    target = await set_target(state, kind)
    await send(state, 'ai', 1)
    assert state.started == []
    await send(state, 'ai', 0, b'first-audio')
    await send(state, 'ai', 0, b'-tail')
    await send(state, 'ai', 2)
    assert state.started == ['ai']
    assert state.ended == [('ai', b'first-audio-tail')]
    assert [c.args[0] for c in state.ws.send_text.call_args_list] == ['START_ACK', 'END_ACK']
    assert await state.targets.get_voice_target(DEVICE) is None
    assert await state.buffers.get_audio_data(target.audio_cache_key) is None
    messages.process_parent_leave_message.assert_not_awaited()
    messages.process_cached_audio.assert_not_awaited()


@pytest.mark.asyncio
@pytest.mark.parametrize('kind', ['parent', 'device'])
@pytest.mark.parametrize('outcome', ['success', 'empty', 'error', 'cancelled'])
async def test_message_terminal_paths_release_only_their_recording(state, kind, outcome):
    target = await set_target(state, kind)
    await send(state, 'message', 1)
    if outcome != 'empty':
        await send(state, 'message', 4, b'message-audio')
    processor = messages.process_parent_leave_message if kind == 'parent' else messages.process_cached_audio
    if outcome == 'error':
        processor.side_effect = RuntimeError('upload failed')
        with pytest.raises(RuntimeError, match='upload failed'):
            await send(state, 'message', 2)
    else:
        await send(state, 'message', 3 if outcome == 'cancelled' else 2)
    assert await state.targets.get_voice_target(DEVICE) is None
    assert await state.buffers.get_audio_data(target.audio_cache_key) is None
    assert state.started == []
    if outcome == 'cancelled':
        processor.assert_not_awaited()
        assert [c.args[0] for c in state.ws.send_text.call_args_list][-2:] == ['INTERRUPT_ACK', 'INTERRUPT_PROCESSED']
    else:
        processor.assert_awaited_once()
    await send(state, 'next-ai', 1)
    await send(state, 'next-ai', 0, b'next-question')
    await send(state, 'next-ai', 2)
    assert state.ended == [('next-ai', b'next-question')]


@pytest.mark.asyncio
async def test_cancelled_message_cannot_receive_late_audio_or_end(state):
    await set_target(state, 'parent')
    await send(state, 'cancelled', 1)
    await send(state, 'cancelled', 4, b'old')
    await send(state, 'cancelled', 3)
    await set_target(state, 'parent')
    fresh = await state.targets.get_voice_target(DEVICE)
    await send(state, 'fresh', 1)
    await send(state, 'fresh', 4, b'new')
    await send(state, 'cancelled', 4, b'late')
    await send(state, 'cancelled', 2)
    assert await state.targets.get_voice_target(DEVICE) is fresh
    assert await state.buffers.get_audio_data(fresh.audio_cache_key) == b'new'
    messages.process_parent_leave_message.assert_not_awaited()


@pytest.mark.asyncio
@pytest.mark.parametrize('claimed', [False, True])
async def test_timeout_clears_pending_or_partial_message_and_allows_ai(state, claimed):
    target = await set_target(state, 'parent')
    if claimed:
        await send(state, 'expired', 1)
        await send(state, 'expired', 4, b'partial')
    state.now[0] += 11
    assert await state.targets.get_voice_target(DEVICE) is None
    assert await state.buffers.get_audio_data(target.audio_cache_key) is None
    await send(state, 'ai', 1)
    await send(state, 'ai', 0, b'question')
    await send(state, 'ai', 2)
    assert state.ended == [('ai', b'question')]


@pytest.mark.asyncio
async def test_active_message_refreshes_idle_timeout(state):
    target = await set_target(state, 'parent')
    await send(state, 'message', 1)
    for _ in range(4):
        state.now[0] += 9
        await send(state, 'message', 4, b'a')
    assert await state.targets.get_voice_target(DEVICE) is target
    assert await state.buffers.get_audio_data(target.audio_cache_key) == b'aaaa'


@pytest.mark.asyncio
async def test_replacement_during_upload_survives_old_recording_cleanup(state):
    old = await set_target(state, 'parent')
    await send(state, 'old', 1)
    await send(state, 'old', 4, b'old-audio')
    replacement = []

    async def upload(*args):
        await set_target(state, 'parent')
        replacement.append(await state.targets.get_voice_target(DEVICE))

    messages.process_parent_leave_message.side_effect = upload
    await send(state, 'old', 2)
    assert await state.targets.get_voice_target(DEVICE) is replacement[0]
    assert replacement[0].audio_cache_key != old.audio_cache_key
    assert await state.buffers.get_audio_data(old.audio_cache_key) is None


@pytest.mark.asyncio
async def test_new_message_command_during_ai_does_not_hijack_its_end(state):
    await send(state, 'ai', 1)
    await send(state, 'ai', 0, b'question')
    target = await set_target(state, 'parent')
    await send(state, 'ai', 2)
    assert state.ended == [('ai', b'question')]
    assert await state.targets.get_voice_target(DEVICE) is target
    messages.process_parent_leave_message.assert_not_awaited()


@pytest.mark.asyncio
async def test_two_senders_to_same_receiver_use_separate_buffers(state):
    await state.targets.set_target(DEVICE, 'receiver')
    await state.targets.set_target('other-sender', 'receiver')
    a = await state.targets.claim_target(DEVICE, state.ws, 'a')
    b = await state.targets.claim_target('other-sender', state.ws, 'b')
    await state.targets.append_audio(DEVICE, a, state.ws, 'a', b'one')
    await state.targets.append_audio('other-sender', b, state.ws, 'b', b'two')
    await state.targets.remove_target(DEVICE, expected=a)
    assert await state.buffers.get_audio_data(b.audio_cache_key) == b'two'


@pytest.mark.asyncio
@pytest.mark.parametrize('claimed', [False, True])
async def test_disconnect_clears_pending_or_partial_target(state, claimed):
    await state.connections.add_connection(DEVICE, state.ws)
    target = await set_target(state, 'parent')
    if claimed:
        await send(state, 'message', 1)
        await send(state, 'message', 4, b'partial')
        session = await state.sessions.get_session((DEVICE, 'message'))
        session.set_interrupted = AsyncMock()
    await cleanup.cleanup_device_sessions(DEVICE, state.ws)
    assert await state.targets.get_voice_target(DEVICE) is None
    assert await state.buffers.get_audio_data(target.audio_cache_key) is None
    assert await state.sessions.get_all_sessions() == []


@pytest.mark.asyncio
async def test_old_socket_cleanup_and_late_packets_preserve_new_connection(state, monkeypatch):
    old = state.ws
    new = SimpleNamespace(send_text=AsyncMock())

    async def authenticate(ws):
        await state.connections.add_connection(DEVICE, old)
        return True, DEVICE, 'LOCAL'

    async def replacement_connects(*args):
        await state.connections.add_connection(DEVICE, new)
        await set_target(state, 'parent')
        await send(state, 'new', 1, ws=new)
        await send(state, 'new', 4, b'new-audio', ws=new)
        await send(state, 'old-late-start', 1, ws=old)
        await send(state, 'old-late-end', 2, ws=old)

    monkeypatch.setattr(endpoint, 'authenticate_websocket', authenticate)
    monkeypatch.setattr(endpoint, 'handle_websocket_messages', replacement_connects)
    await endpoint.websocket_endpoint(old)
    target = await state.targets.get_voice_target(DEVICE)
    assert target.websocket is new
    assert await state.buffers.get_audio_data(target.audio_cache_key) == b'new-audio'
    assert await state.connections.get_connection(DEVICE) is new
    assert await state.sessions.get_session((DEVICE, 'new')) is not None
    cleanup.task_manager.cancel_session_tasks.assert_not_awaited()


@pytest.mark.asyncio
async def test_command_before_connect_is_adopted_and_cleared_on_early_disconnect(state):
    target = await set_target(state, 'parent')
    assert target.websocket is None
    await state.connections.add_connection(DEVICE, state.ws)
    await state.targets.bind_connection(DEVICE, state.ws)
    await cleanup.cleanup_device_sessions(DEVICE, state.ws)
    assert await state.targets.get_voice_target(DEVICE) is None


@pytest.mark.asyncio
async def test_reconnect_discards_old_partial_message(state):
    await state.connections.add_connection(DEVICE, state.ws)
    old = await set_target(state, 'parent')
    await send(state, 'old', 1)
    await send(state, 'old', 4, b'old')
    new = SimpleNamespace(send_text=AsyncMock())
    await state.connections.add_connection(DEVICE, new)
    await state.targets.bind_connection(DEVICE, new)
    assert await state.targets.get_voice_target(DEVICE) is None
    assert await state.buffers.get_audio_data(old.audio_cache_key) is None
    await send(state, 'new-ai', 1, ws=new)
    await send(state, 'new-ai', 0, b'reconnected-question', ws=new)
    await send(state, 'new-ai', 2, ws=new)
    assert state.ended == [('new-ai', b'reconnected-question')]


@pytest.mark.asyncio
@pytest.mark.parametrize('outcome', ['success', 'empty', 'upload_error'])
async def test_real_parent_message_processing_keeps_notifications_and_clears_state(state, monkeypatch, outcome):
    monkeypatch.setattr(messages, 'process_parent_leave_message', PROCESS_PARENT_MESSAGE)
    upload = AsyncMock(return_value=SimpleNamespace(file_key='local-test/audio.wav'))
    create = AsyncMock(return_value=(SimpleNamespace(child_name='local-test'), None, []))
    notify = Mock()
    if outcome == 'upload_error':
        upload.side_effect = RuntimeError('simulated upload failure')
    monkeypatch.setattr(messages, 'upload_message_audio', upload)
    monkeypatch.setattr(messages.im_conversation_service, 'create_device_parent_leave_message', create)
    monkeypatch.setattr(messages.wechat_mp_notification_service, 'schedule_leave_message_notification', notify)
    await state.connections.add_connection(DEVICE, state.ws)
    target = await set_target(state, 'parent')
    await send(state, 'parent', 1)
    if outcome != 'empty':
        await send(state, 'parent', 4, b'\x01\x00'*400)
    await send(state, 'parent', 2)
    assert await state.targets.get_voice_target(DEVICE) is None
    assert await state.buffers.get_audio_data(target.audio_cache_key) is None
    if outcome == 'success':
        upload.assert_awaited_once()
        create.assert_awaited_once()
        assert create.call_args.kwargs['audio_content'].startswith(b'RIFF')
        notify.assert_called_once()
        assert state.ws.send_text.call_args.args[0].startswith('PROMPT_SOUND_URL:')
    else:
        create.assert_not_awaited()
        notify.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.parametrize('first_kind', [0, 4])
async def test_recording_cannot_switch_between_ai_and_message_after_audio(state, first_kind):
    target = await set_target(state, 'parent')
    await send(state, 'mixed', 1)
    await send(state, 'mixed', first_kind, b'first')
    await send(state, 'mixed', 4 if first_kind == 0 else 0, b'wrong-mode')
    if first_kind == 0:
        await send(state, 'mixed', 2)
        assert state.ended == [('mixed', b'first')]
        messages.process_parent_leave_message.assert_not_awaited()
    else:
        assert await state.buffers.get_audio_data(target.audio_cache_key) == b'first'
        assert state.started == []
        await send(state, 'mixed', 2)
        messages.process_parent_leave_message.assert_awaited_once()


@pytest.mark.asyncio
async def test_new_recording_recovers_abandoned_message_without_reusing_audio(state):
    abandoned = await set_target(state, 'parent')
    await send(state, 'abandoned', 1)
    await send(state, 'abandoned', 4, b'abandoned-audio')
    await send(state, 'new-ai', 1)
    await send(state, 'new-ai', 0, b'question')
    await send(state, 'abandoned', 2, active=(DEVICE, 'new-ai'))
    await send(state, 'new-ai', 2)
    assert state.ended == [('new-ai', b'question')]
    assert await state.buffers.get_audio_data(abandoned.audio_cache_key) is None
    messages.process_parent_leave_message.assert_not_awaited()