#!/usr/bin/env python3
"""Deterministic Agent + PCM downlink over local authenticated TLS WebSocket.

No provider calls, production service, device credentials or hardware. Reports
frame metadata and hashes, never PCM, base64, tokens or session histories.
"""
import argparse
import asyncio
import base64
import hashlib
import json
import math
import os
from pathlib import Path
import socket
import ssl
import struct
import subprocess
import sys
import tempfile
import time

import httpx
import uvicorn
from websockets.asyncio.client import connect

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO/'server'))
sys.path.insert(0, str(REPO/'server/tests'))
from fake_device import DownlinkRecorder


def deterministic_pcm(seconds):
    # Exactly seconds * 88200 little-endian mono bytes, nonzero and reproducible.
    period = b''.join(struct.pack('<h', round(5000*math.sin(2*math.pi*i/100))) for i in range(100))
    size = seconds*44100*2
    return (period*((size+len(period)-1)//len(period)))[:size]


class ContinuousDevice:
    def __init__(self):
        self.recorder = DownlinkRecorder()
        self.events = asyncio.Queue()
        self.playback = asyncio.Queue(maxsize=8)
        self.queue_high_water = 0
        self.reader = None

    async def consume(self):
        while True:
            size = await self.playback.get()
            await asyncio.sleep(size/88200)
            self.playback.task_done()

    async def read(self, ws):
        try:
            async for wire in ws:
                event = self.recorder.receive(wire)
                if event['type'] == 'speech_chunk':
                    size = len(base64.b64decode(event['audio'], validate=True))
                    self.playback.put_nowait(size)
                    self.queue_high_water = max(self.queue_high_water, self.playback.qsize())
                else:
                    self.events.put_nowait(event)
        finally:
            self.recorder.disconnect(ws.close_code, expected=ws.close_code == 1000)

    async def wait(self, kind, **fields):
        async with asyncio.timeout(35):
            while True:
                take = asyncio.create_task(self.events.get())
                try:
                    done, _ = await asyncio.wait([take, self.reader], return_when=asyncio.FIRST_COMPLETED)
                    if self.reader in done:
                        await self.reader
                        raise AssertionError('receiver disconnected before requested event')
                    event = take.result()
                finally:
                    if not take.done():
                        take.cancel()
                        await asyncio.gather(take, return_exceptions=True)
                assert event['type'] not in ('speech_cancel', 'speech_error'), event['type']
                if event['type'] == kind and all(event.get(k) == v for k, v in fields.items()):
                    return event


async def main(output):
    with tempfile.TemporaryDirectory(prefix='downlink-wss-') as folder:
        root = Path(folder)
        credentials = root/'credentials.json'
        credentials.write_text(json.dumps(dict(device_id='downlink-fake', device_token='test-only', pi_token='pi-test-only')))
        credentials.chmod(0o600)
        os.environ.update(DEVICE_CREDENTIALS_FILE=str(credentials), GATEWAY_STATE_FILE=str(root/'targets.json'),
                          GATEWAY_PROJECTS_FILE=str(root/'projects.json'), GATEWAY_SPEECH_FILE=str(root/'speech.json'),
                          DEVICE_LAB_CONTROL_DIR=str(root/'runtime'), GATEWAY_CODEX_ENABLED='0')
        import device_gateway.app as gateway
        class PCMProvider:
            pcm = b''
            async def synthesize(self, text, language): return self.pcm
        provider = PCMProvider()
        gateway.adapters = {}  # Registry fallback, no Pi or Codex subprocesses.
        gateway.speeches.provider = provider
        send_times = {}
        original_send = gateway.speeches.send
        async def observed_send(device, payload):
            if payload['type'] == 'speech_chunk':
                send_times.setdefault(payload['speech_id'], []).append(time.monotonic())
            await original_send(device, payload)
        gateway.speeches.send = observed_send
        gateway.snapshot_interval_seconds = .25  # Concurrent state snapshots while receiving audio.
        cert, key = root/'cert.pem', root/'key.pem'
        subprocess.run(['openssl', 'req', '-x509', '-newkey', 'rsa:2048', '-nodes', '-days', '1',
                        '-subj', '/CN=localhost', '-addext', 'subjectAltName=DNS:localhost,IP:127.0.0.1',
                        '-keyout', str(key), '-out', str(cert)], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        context = ssl.create_default_context(cafile=str(cert))
        sock = socket.socket()
        sock.bind(('127.0.0.1', 0))
        port = sock.getsockname()[1]
        server = uvicorn.Server(uvicorn.Config(gateway.app, log_level='error', ssl_certfile=str(cert),
                                             ssl_keyfile=str(key), ws_max_size=8192, ws_per_message_deflate=False))
        serving = asyncio.create_task(server.serve(sockets=[sock]))
        device = ContinuousDevice()
        consumer = asyncio.create_task(device.consume())
        report = dict(scope='SOFTWARE ONLY: local TLS WSS, deterministic PCM, modeled realtime consumer',
                      cases=[], agent_downlink={})
        try:
            async with asyncio.timeout(10):
                while not server.started:
                    if serving.done(): await serving; raise AssertionError('test server did not start')
                    await asyncio.sleep(.01)
            async with httpx.AsyncClient(base_url=f'https://127.0.0.1:{port}', verify=context) as client:
                async with connect(f'wss://127.0.0.1:{port}/device/ws', ssl=context, compression=None) as ws:
                    device.reader = asyncio.create_task(device.read(ws))
                    await ws.send(json.dumps(dict(type='hello', device_id='downlink-fake', device_token='test-only',
                        board='fake-device', firmware='0.5.1', capabilities=['agent_target', 'speech_playback'])))
                    await device.wait('hello_ack'); await device.wait('speech_settings')
                    async def pi(kind, **extra):
                        response = await client.post('/integrations/pi/events', headers={'x-device-lab-pi-token':'pi-test-only'},
                            json=dict(event=kind, session_id='s', name='Public test', cwd=folder, turn_id='t', **extra))
                        response.raise_for_status()
                    await pi('session_start')
                    selected = await client.post('/devices/downlink-fake/target', json={'target':'pi:s'})
                    selected.raise_for_status()
                    assert selected.json()['target'] == 'agent:pi:s'
                    await device.wait('agent_state', state='IDLE')
                    await pi('agent_start'); await device.wait('agent_state', state='RUNNING')
                    await pi('activity', activity_type='tool', activity='Running\nsleep 3')
                    assert (await device.wait('agent_activity'))['activity'] == 'Running\nsleep 3'
                    await pi('activity', activity_type='text', activity='正在检查。')
                    assert (await device.wait('agent_activity'))['activity'] == '正在检查。'
                    await pi('agent_end', final_text='检查完成。')
                    assert (await device.wait('agent_state', state='DONE'))['final_text'] == '检查完成。'
                    report['agent_downlink'] = dict(running=True, tool=True, text=True, final=True)
                    for seconds in (3, 10, 20):
                        provider.pcm = deterministic_pcm(seconds)
                        device.queue_high_water = 0
                        await ws.send(json.dumps(dict(type='speech_request', request_id=f'pcm-{seconds}',
                            target='agent:pi:s', source='screen', text='Deterministic transport test.')))
                        start = await device.wait('speech_start')
                        await device.wait('speech_end', speech_id=start['speech_id'])
                        await asyncio.wait_for(device.playback.join(), 3)
                        result = device.recorder.report(start['speech_id'])
                        times = send_times[start['speech_id']]
                        send_intervals = [b-a for a, b in zip(times, times[1:])]
                        result['minimum_send_interval_seconds'] = min(send_intervals)
                        result['maximum_send_interval_seconds'] = max(send_intervals)
                        assert result['raw_pcm_bytes'] == seconds*88200
                        assert result['pcm_sha256'] == hashlib.sha256(provider.pcm).hexdigest()
                        assert result['chunk_count'] == (len(provider.pcm)+7055)//7056
                        assert result['maximum_websocket_text_payload_bytes'] < 10000
                        # TLS/TCP may coalesce receive deliveries. Verify pacing at
                        # the sender, and receive burst tolerance with the bounded queue.
                        assert result['minimum_send_interval_seconds'] >= .079
                        assert seconds*.98 <= result['total_transfer_seconds'] < seconds*1.35+1
                        assert device.queue_high_water <= 8
                        assert not device.recorder.disconnects and gateway.devices['downlink-fake'].online
                        result.update(audio_seconds=seconds, queue_high_water=device.queue_high_water,
                                      missing_sequences=0, duplicate_sequences=0, server_disconnects=0)
                        report['cases'].append(result)
                        print(json.dumps(result, ensure_ascii=False), flush=True)
                await device.reader  # Record expected client close separately from server failures.
        finally:
            if device.reader and not device.reader.done(): device.reader.cancel()
            consumer.cancel()
            await asyncio.gather(consumer, *([device.reader] if device.reader else []), return_exceptions=True)
            server.should_exit = True
            await asyncio.wait_for(serving, 10)
            sock.close()
            report['records'] = device.recorder.records
            report['disconnects'] = device.recorder.disconnects
            output.parent.mkdir(parents=True, exist_ok=True)
            output.write_text(json.dumps(report, ensure_ascii=False, indent=2)+'\n')
        assert len(report['cases']) == 3
        print('Fake-device downlink PASS; physical receive/playback NOT tested.', flush=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('/tmp/downlink-transport.json'))
    args = parser.parse_args()
    asyncio.run(main(args.output))