#!/usr/bin/env python3
"""
Simulate one audio Node for the ESP32 Hub over the node UDP protocol.

The Hub (see src/network/node_protocol.h) runs:
  * control plane on UDP 4210 (PROBE / START / STOP / POLL)
  * data plane   on UDP 4211 (DATA packets from Nodes)

This script behaves like a Node with 4 microphones:
  * answers PROBE with a PROBE_RESP (so the Hub registers it)
  * streams 4 interleaved 48 kHz / 16-bit sine tones (one per mic) as
    DATA packets while recording, paced to real time so the WAV plays back
    at the correct pitch and duration
  * sends HEARTBEATs so the Hub keeps it marked online

Usage:
    # PC connected to the Hub's AP (192.168.4.x), or any host on the subnet
    python3 scripts/simulate_node.py

    # multiple virtual nodes (distinct tones):
    python3 scripts/simulate_node.py --node-id 2 --frequencies 330,415,495,660

The four default tones are an A-major chord, one per mic, so each channel is
clearly distinguishable when the 40-channel WAV is played back.

Requires only the Python standard library.
"""

import argparse
import math
import select
import signal
import socket
import struct
import sys
import time

# --- Protocol constants (mirror src/network/network_config.h + node_protocol.h)

PROTOCOL_VERSION = 1
CONTROL_PORT = 4210
DATA_PORT = 4211

MSG_PROBE = 0x01
MSG_PROBE_RESP = 0x02
MSG_START = 0x03
MSG_STOP = 0x04
MSG_HEARTBEAT = 0x05
MSG_POLL = 0x06
MSG_POLL_ACK = 0x07
MSG_DATA = 0x10

STATUS_IDLE = 0
STATUS_RECORDING = 1

MICS_PER_NODE = 1
SAMPLE_RATE = 16000
BITS_PER_SAMPLE = 16
FRAME_BYTES = MICS_PER_NODE * 2          # 8 bytes per frame (4 ch x 2 B)

# Frames per DATA packet: (header 8 B) + n * 8 B must stay below 1472.
FRAMES_PER_PACKET = 100

AMPLITUDE = 8000.0                        # ~0.25 x full scale
HEARTBEAT_MS = 2000


def u16(x):
    return struct.pack('<H', x & 0xFFFF)


def u32(x):
    return struct.pack('<I', x & 0xFFFFFFFF)


def build_probe_resp(node_id, mac, firmware, mic_count, sample_rate,
                     bits_per_sample, buffer_bytes, status):
    fw = firmware.encode('ascii')[:15].ljust(15, b'\x00') + b'\x00'
    return (bytes([MSG_PROBE_RESP, PROTOCOL_VERSION, node_id])
            + mac
            + fw
            + bytes([mic_count])
            + u32(sample_rate)
            + u16(bits_per_sample)
            + u32(buffer_bytes)
            + bytes([status]))


def build_heartbeat(node_id, uptime_ms, buffer_fill, status):
    return (bytes([MSG_HEARTBEAT, PROTOCOL_VERSION, node_id])
            + u32(uptime_ms)
            + u32(buffer_fill)
            + bytes([status]))


def build_data(node_id, mic_count, frame_count, pcm):
    return (bytes([MSG_DATA, PROTOCOL_VERSION, node_id, mic_count])
            + u32(frame_count)
            + pcm)


class SineGenerator:
    """Four independent phase accumulators producing interleaved 16-bit PCM."""

    def __init__(self, frequencies, amplitude):
        self.frequencies = frequencies
        self.amplitude = amplitude
        self.phase = [0.0] * len(frequencies)
        self._two_pi = 2.0 * math.pi

    def frames(self, count):
        """Generate `count` interleaved frames (count * FRAME_BYTES bytes)."""
        pcm = bytearray(count * FRAME_BYTES)
        step = [self._two_pi * f / SAMPLE_RATE for f in self.frequencies]

        for f in range(count):
            for ch in range(len(self.frequencies)):
                self.phase[ch] += step[ch]
                if self.phase[ch] >= self._two_pi:
                    self.phase[ch] -= self._two_pi
                sample = int(self.amplitude * math.sin(self.phase[ch]))
                off = (f * MICS_PER_NODE + ch) * 2
                pcm[off] = sample & 0xFF
                pcm[off + 1] = (sample >> 8) & 0xFF

        return bytes(pcm)


class SimulatedNode:
    def __init__(self, args):
        self.node_id = args.node_id
        self.mac = bytes([0x02, 0x00, 0x00, 0x00, 0x00, 0x10 + (args.node_id & 0xEF)])
        self.firmware = "py-node-v1"
        self.buffer_bytes = 64 * 1024

        self.hub_ip = args.hub_ip
        self.recording = False
        self.status = STATUS_IDLE
        self.uptime_ms = 0
        self.started_at = 0.0
        self.produced = 0
        self.last_heartbeat = 0.0
        self.running = True

        self.sine = SineGenerator(args.frequencies, args.amplitude)
        self.stop_after = args.duration

        self.ctrl = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.ctrl.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            self.ctrl.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        except OSError:
            pass
        self.ctrl.bind(('', CONTROL_PORT))

        self.data = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            self.data.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        except OSError:
            pass

    # --- helpers ----------------------------------------------------------

    def send_probe_resp(self):
        pkt = build_probe_resp(self.node_id, self.mac, self.firmware,
                               MICS_PER_NODE, SAMPLE_RATE, BITS_PER_SAMPLE,
                               self.buffer_bytes, self.status)
        self.ctrl.sendto(pkt, (self.hub_ip, CONTROL_PORT))
        print(f"[node {self.node_id}] PROBE_RESP sent to {self.hub_ip}")

    def send_heartbeat(self):
        pkt = build_heartbeat(self.node_id, self.uptime_ms,
                              self.buffered_frames() * FRAME_BYTES, self.status)
        self.ctrl.sendto(pkt, (self.hub_ip, CONTROL_PORT))

    def buffered_frames(self):
        if not self.recording:
            return 0
        return max(0, int(time.monotonic() - self.started_at) * SAMPLE_RATE
                   - self.produced)

    def stream_pacing(self):
        """Send the frames that real time says are due, in DATA packets."""
        if not self.recording or self.hub_ip is None:
            return

        target = int((time.monotonic() - self.started_at) * SAMPLE_RATE)
        while self.produced < target:
            n = min(target - self.produced, FRAMES_PER_PACKET)
            pcm = self.sine.frames(n)
            pkt = build_data(self.node_id, MICS_PER_NODE, n, pcm)
            self.data.sendto(pkt, (self.hub_ip, DATA_PORT))
            self.produced += n

    # --- control packet handling ------------------------------------------

    def handle(self, packet, src):
        if len(packet) < 2:
            return
        msg, ver = packet[0], packet[1]
        if ver != PROTOCOL_VERSION:
            return

        if src[0] != self.hub_ip:
            self.hub_ip = src[0]          # learn the Hub from any control packet
            print(f"[node {self.node_id}] learned Hub at {self.hub_ip}")

        if msg == MSG_PROBE:
            self.send_probe_resp()

        elif msg == MSG_START:
            self.recording = True
            self.status = STATUS_RECORDING
            self.started_at = time.monotonic()
            self.produced = 0
            print(f"[node {self.node_id}] START - streaming 4 tones at "
                  f"{SAMPLE_RATE} Hz, {MICS_PER_NODE} ch")

        elif msg == MSG_STOP:
            if self.recording:
                print(f"[node {self.node_id}] STOP - sent "
                      f"{self.produced} frames")
            self.recording = False
            self.status = STATUS_IDLE

        elif msg == MSG_POLL:
            # We stream continuously; nothing extra to send.
            pass

    # --- main loop --------------------------------------------------------

    def run(self):
        print(f"[node {self.node_id}] listening on UDP {CONTROL_PORT}...")
        print(f"[node {self.node_id}] 4 tones: {self.sine.frequencies} Hz")
        if self.hub_ip:
            print(f"[node {self.node_id}] Hub set to {self.hub_ip}")
        else:
            print("[node %d] waiting for a PROBE to learn the Hub IP "
                  "(PC must be on the Hub's AP network)" % self.node_id)

        while self.running:
            ready, _, _ = select.select([self.ctrl], [], [], 0.002)
            if ready:
                try:
                    packet, src = self.ctrl.recvfrom(2048)
                    self.handle(packet, src)
                except OSError:
                    pass

            if self.recording:
                self.stream_pacing()
                if self.stop_after is not None:
                    if time.monotonic() - self.started_at >= self.stop_after:
                        print(f"[node {self.node_id}] --duration reached, "
                              "waiting for STOP")
                        self.recording = False
                        self.status = STATUS_IDLE

            now_ms = int(time.monotonic() * 1000)
            self.uptime_ms = now_ms
            if self.hub_ip and now_ms - self.last_heartbeat >= HEARTBEAT_MS:
                self.last_heartbeat = now_ms
                self.send_heartbeat()

        self.ctrl.close()
        self.data.close()


def parse_args(argv):
    parser = argparse.ArgumentParser(
        description='Simulate a 4-mic audio Node for the ESP32 Hub.')
    parser.add_argument('--node-id', type=int, default=1,
                        help='Node ID (default 1). Overrides the matching '
                             'built-in simulator node.')
    parser.add_argument('--hub-ip', default=None,
                        help='Hub IP, e.g. 192.168.4.1 (auto-learned from '
                             'the first control packet if omitted).')
    parser.add_argument('--frequencies', default='440,554.37,659.25,880',
                        help='Comma-separated sine frequencies per mic '
                             '(default an A-major chord).')
    parser.add_argument('--amplitude', type=float, default=AMPLITUDE,
                        help='Sine amplitude 0..32767 (default 8000).')
    parser.add_argument('--duration', type=float, default=None,
                        help='Auto-stop after N seconds of recording '
                             '(default: until the Hub sends STOP).')
    return parser.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    args.frequencies = [float(f) for f in args.frequencies.split(',')]
    if len(args.frequencies) < MICS_PER_NODE:
        print('error: provide at least %d frequencies' % MICS_PER_NODE,
              file=sys.stderr)
        return 1

    node = SimulatedNode(args)

    def on_sigint(_sig, _frame):
        node.running = False
        print('\n[node %d] shutting down' % args.node_id)

    signal.signal(signal.SIGINT, on_sigint)
    node.run()
    return 0


if __name__ == '__main__':
    sys.exit(main())