"""Physical EEG acquisition bridge. Requires BrainFlow and websockets 16+.
No synthetic board, recording replay, stimulation or model inference is provided.
"""
import argparse
import asyncio
import hmac
import json
import math
import os
import ssl


def prepare_frame(data, rows, timestamp_row, last_timestamp, sample_rate):
    """Validate acquisition timestamps before sending nominal-rate EEG samples."""
    times = data[timestamp_row].tolist()
    for timestamp in times:
        if not math.isfinite(timestamp):
            raise ValueError("Acquisition timestamp is non-finite")
        if last_timestamp is not None:
            elapsed = timestamp - last_timestamp
            if elapsed <= 0 or elapsed > 1.5 / sample_rate:
                raise ValueError("Acquisition timing gap or reversal; reconnect and recalibrate")
        last_timestamp = timestamp
    samples = [data[row].tolist() for row in rows]
    if any(not math.isfinite(v) or abs(v) > 1e7 for channel in samples for v in channel):
        raise ValueError("Invalid channel voltage")
    return samples, last_timestamp


async def serve_board(args, board_class, params_class, serve):
    token = os.environ.get('NOETRYX_STREAM_TOKEN', '')
    if not 24 <= len(token) <= 256:
        raise ValueError('Set NOETRYX_STREAM_TOKEN to a random secret of 24–256 characters.')
    if args.board_id < 0 or not args.confirm_microvolts:
        raise ValueError('Choose a physical board and confirm EEG microvolt units. Synthetic/replay boards are rejected.')
    rate = board_class.get_sampling_rate(args.board_id)
    rows = list(board_class.get_eeg_channels(args.board_id))
    if not 128 <= rate <= 1024 or not 1 <= len(rows) <= 16:
        raise ValueError('This bridge supports 128–1024 Hz and 1–16 EEG channels.')
    timestamp_row = board_class.get_timestamp_channel(args.board_id)
    descriptor = board_class.get_board_descr(args.board_id)
    names = descriptor.get('eeg_names', '').split(',')
    if len(names) != len(rows) or any(not n for n in names):
        names = ['EEG ' + str(i + 1) for i in range(len(rows))]
    params = params_class()
    params.serial_port = args.serial_port
    params.mac_address = args.mac_address
    if args.params_json:
        with open(args.params_json, encoding='utf-8') as f:
            for key, value in json.load(f).items():
                if key.startswith('_') or not hasattr(params, key):
                    raise ValueError('Unsupported BrainFlow parameter: ' + key)
                setattr(params, key, value)
    tls = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    tls.minimum_version = ssl.TLSVersion.TLSv1_2
    tls.load_cert_chain(args.cert, args.key)
    lock = asyncio.Lock()

    async def handler(ws):
        try:
            message = json.loads(await asyncio.wait_for(ws.recv(), timeout=5))
            supplied = message.get('token', '')
            if message.get('type') != 'authorize' or not isinstance(supplied, str) or not hmac.compare_digest(supplied, token):
                await ws.close(code=1008, reason='Authorization required')
                return
        except Exception:
            await ws.close(code=1008, reason='Authorization required')
            return
        if lock.locked():
            await ws.close(code=1013, reason='A receiver is already connected')
            return
        async with lock:
            board = board_class(args.board_id, params)
            prepared = streaming = False
            try:
                board.prepare_session()
                prepared = True
                board.start_stream()
                streaming = True
                await ws.send(json.dumps({'type': 'hello', 'protocol': 'noetryx-eeg-1',
                    'source': 'device', 'sourceName': descriptor.get('name', 'BrainFlow device'),
                    'sampleRate': rate, 'channels': names, 'unit': 'uV'}))
                sequence = first_sample = 0
                last_timestamp = None
                while first_sample < rate * 600:
                    await asyncio.sleep(0.25)
                    data = board.get_board_data()
                    count = data.shape[1]
                    if not count:
                        continue
                    if count > rate * 2:
                        raise ValueError('Receiver fell behind acquisition; reconnect')
                    samples, last_timestamp = prepare_frame(data, rows, timestamp_row, last_timestamp, rate)
                    await ws.send(json.dumps({'type': 'samples', 'sequence': sequence,
                        'firstSample': first_sample, 'samples': samples}, allow_nan=False))
                    sequence += 1
                    first_sample += count
                await ws.close(code=1000, reason='Ten-minute acquisition completed')
            except Exception as error:
                print('Acquisition stopped:', type(error).__name__, str(error))
                await ws.close(code=1011, reason='Acquisition stopped; inspect bridge console')
            finally:
                if streaming:
                    try:
                        board.stop_stream()
                    finally:
                        if prepared:
                            board.release_session()
                elif prepared:
                    board.release_session()

    async with serve(handler, args.host, args.port, ssl=tls, origins=[args.origin],
                     max_size=1024, max_queue=4, ping_interval=20, ping_timeout=20):
        print('EEG bridge ready. Keep this computer running. Only the configured origin and access token are accepted.')
        await asyncio.get_running_loop().create_future()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--board-id', type=int, required=True)
    parser.add_argument('--serial-port', default='')
    parser.add_argument('--mac-address', default='')
    parser.add_argument('--params-json')
    parser.add_argument('--confirm-microvolts', action='store_true', required=True)
    parser.add_argument('--cert', required=True)
    parser.add_argument('--key', required=True)
    parser.add_argument('--host', default='127.0.0.1')
    parser.add_argument('--port', type=int, default=8765)
    parser.add_argument('--origin', default='https://noetryx.earthy-cloud-7779.chatgpt.site')
    args = parser.parse_args()
    from brainflow.board_shim import BoardShim, BrainFlowInputParams
    from websockets.asyncio.server import serve
    asyncio.run(serve_board(args, BoardShim, BrainFlowInputParams, serve))


if __name__ == '__main__':
    main()