"""Test-only device receiver. Traces contain metadata, never tokens or PCM/base64."""
import base64
from datetime import datetime, timezone
import hashlib
import json
import time


class DownlinkRecorder:
    def __init__(self):
        self.started = time.monotonic()
        self.records = []
        self.streams = {}
        self.disconnects = []

    def receive(self, wire):
        assert isinstance(wire, str), 'Gateway must send WebSocket TEXT'
        event = json.loads(wire)
        now = time.monotonic()
        kind = event['type']
        record = dict(order=len(self.records), timestamp=datetime.now(timezone.utc).isoformat(),
                      elapsed_seconds=now-self.started, type=kind, payload_bytes=len(wire.encode('utf-8')))
        for key in ('speech_id', 'sequence', 'chunks', 'target', 'session_id', 'agent', 'state', 'activity_type'):
            if key in event:
                record[key] = event[key]
        for key in ('activity', 'recent_text', 'final_text'):
            if key in event:
                record[key + '_bytes'] = len(event[key].encode('utf-8'))
        self.records.append(record)
        if kind == 'speech_start':
            sid = event['speech_id']
            assert sid not in self.streams
            assert (event['sample_rate'], event['bits'], event['channels']) == (44100, 16, 1)
            self.streams[sid] = dict(start=now, times=[], raw_bytes=0, sizes=[], digest=hashlib.sha256(), end=None)
        elif kind == 'speech_chunk':
            stream = self.streams[event['speech_id']]
            assert stream['end'] is None
            assert event['sequence'] == len(stream['times']), 'missing or duplicate speech sequence'
            pcm = base64.b64decode(event['audio'], validate=True)
            assert 0 < len(pcm) <= 7056 and len(pcm) % 2 == 0
            stream['times'].append(now)
            stream['sizes'].append(len(pcm))
            stream['raw_bytes'] += len(pcm)
            stream['digest'].update(pcm)
        elif kind == 'speech_end':
            stream = self.streams[event['speech_id']]
            assert stream['end'] is None and event['chunks'] == len(stream['times'])
            stream['end'] = now
        return event

    def disconnect(self, code, *, expected=False):
        self.disconnects.append(dict(timestamp=datetime.now(timezone.utc).isoformat(),
                                     elapsed_seconds=time.monotonic()-self.started, code=code, expected=expected))

    def report(self, sid):
        stream = self.streams[sid]
        assert stream['end'] is not None
        duration = stream['end'] - stream['start']
        intervals = [b-a for a, b in zip(stream['times'], stream['times'][1:])]
        frames = [r for r in self.records if r.get('speech_id') == sid]
        return dict(raw_pcm_bytes=stream['raw_bytes'], chunk_count=len(stream['times']),
                    pcm_sha256=stream['digest'].hexdigest(),
                    maximum_websocket_text_payload_bytes=max(r['payload_bytes'] for r in frames),
                    total_transfer_seconds=duration, average_pcm_bytes_per_second=stream['raw_bytes']/duration,
                    average_text_bytes_per_second=sum(r['payload_bytes'] for r in frames)/duration,
                    minimum_chunk_interval_seconds=min(intervals, default=0),
                    maximum_chunk_interval_seconds=max(intervals, default=0))


class FakeDevice:
    """In-process FastAPI websocket client using the same recorder as the TLS smoke."""
    def __init__(self, ws, device_id='downlink-test', token='test-device-only'):
        self.ws, self.recorder = ws, DownlinkRecorder()
        ws.send_json(dict(type='hello', device_id=device_id, device_token=token,
                          board='fake-device', firmware='0.5.1',
                          capabilities=['agent_target', 'speech_playback']))
        self.wait('hello_ack')

    def wait(self, kind, **fields):
        for _ in range(300):
            message = self.ws.receive()
            if message['type'] == 'websocket.close':
                self.recorder.disconnect(message['code'])
                raise AssertionError(f"unexpected server disconnect: {message['code']}")
            event = self.recorder.receive(message['text'])
            if event['type'] == kind and all(event.get(k) == v for k, v in fields.items()):
                return event
        raise AssertionError(f'missing {kind}')