#!/usr/bin/env python3
"""Record firmware telemetry continuously (requires pyserial; no UDP subscriber)."""
import argparse
from datetime import datetime, timezone
import json
import math
from pathlib import Path
import re
import sys
import time

import serial


FIELDS = {
    'submitted_fps': (r'([\d.]+) fps submitted', float),
    'attempted_fps': (r'([\d.]+) fps attempted', float),
    'kilobytes_per_second': (r'([\d.]+) kB/s', float),
    'send_failures': (r'(\d+) send failures', int),
    'retries': (r'(\d+) retries', int),
    'peer_frames_dropped': (r'(\d+) peer frames dropped', int),
    'send_ms': (r'([\d.]+) ms/send', float),
    'subscribers': (r'(\d+) clients subscribed', int),
    'heap_free': (r'internal heap (\d+) free', int),
    'heap_largest': (r'free / (\d+) largest bytes', int),
    'rssi_dbm': (r'RSSI (-?\d+) dBm', int),
    'max_tx_power_dbm': (r'max TX power ([\w.+-]+) dBm', float),
    'chip_temperature_c': (r'chip temp ([\w.+-]+) C', float),
    'target_fps': (r'target (\d+) fps', int),
    'capture_errors': (r'capture errors (\d+)', int),
    'exposure_lines': (r'exposure (-?\d+) lines', int),
    'jpeg_quality': (r'jpeg quality (\d+)', int),
    'gain_multiplier': (r'gain ([\d.]+)x', float),
    'capture_queue_skipped': (r'capture queue skipped (\d+)', int),
    'stale_frames_dropped': (r'stale frames dropped (\d+)', int),
    'deadline_peer_drops': (r'deadline peer drops (\d+)', int),
    'max_submitted_capture_age_ms': (r'max submitted capture age ([\d.]+) ms', float),
    'max_send_hold_ms': (r'max send hold ([\d.]+) ms', float),
}


def parse_sample(line):
    if 'fps submitted' not in line:
        return None
    sample = {}
    match = re.search(r'\((\d+)\)', line)
    if match:
        sample['uptime_ms'] = int(match[1])
    for name, (pattern, convert) in FIELDS.items():
        match = re.search(pattern, line)
        if match:
            value = convert(match[1])
            sample[name] = value if not isinstance(value, float) or math.isfinite(value) else None
    sample['thermal_limited'] = 'thermal_limited true' in line
    return sample


def write_json(path, data):
    temporary = path.with_suffix(path.suffix+'.tmp')
    temporary.write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n', encoding='utf-8')
    temporary.replace(path)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--port', default='COM7')
    parser.add_argument('--seconds', type=float, default=1860)
    parser.add_argument('--out', type=Path, required=True, help='Output path prefix')
    parser.add_argument('--reset', action='store_true')
    parser.add_argument('--stop-file', type=Path, help='Finish cleanly when this file exists')
    args = parser.parse_args()
    args.out.parent.mkdir(parents=True, exist_ok=True)
    port = serial.Serial()
    port.port, port.baudrate, port.timeout = args.port, 115200, .2
    port.dtr = port.rts = False
    records, events = [], []
    started = time.perf_counter()
    last_uptime, reboot_count, last_sample_at = None, 0, None
    serial_error = None
    try:
        port.open()
        if args.reset:
            port.rts = True
            time.sleep(.15)
            port.rts = False
        with args.out.with_suffix('.serial.log').open('w', encoding='utf-8') as raw, args.out.with_suffix('.telemetry.ndjson').open('w', encoding='utf-8') as telemetry:
            while time.perf_counter()-started < args.seconds:
                if args.stop_file and args.stop_file.exists():
                    break
                line = port.readline().decode('utf-8', errors='replace')
                if not line:
                    continue
                raw.write(line)
                raw.flush()
                elapsed = time.perf_counter()-started
                sample = parse_sample(line)
                if sample:
                    sample.update(elapsed_seconds=round(elapsed, 3), utc=datetime.now(timezone.utc).isoformat())
                    uptime = sample.get('uptime_ms')
                    if uptime is not None and last_uptime is not None and uptime+1000 < last_uptime:
                        reboot_count += 1
                    last_uptime = uptime
                    last_sample_at = elapsed
                    records.append(sample)
                    telemetry.write(json.dumps(sample)+'\n')
                    telemetry.flush()
                    write_json(args.out.with_suffix('.latest.json'), dict(elapsed_seconds=round(elapsed, 3), samples=len(records), reboot_count=reboot_count, latest=sample))
                    if len(records) % 12 == 0:
                        print(json.dumps(dict(minutes=round(elapsed/60, 1), reboot_count=reboot_count, latest=sample)), flush=True)
                elif re.search(r'panic|Guru Meditation|Brownout|watchdog|NO-EOI|OVF|thermal load limit|reset reason|capture cap|configured capture', line, re.I):
                    events.append(dict(elapsed_seconds=round(elapsed, 3), text=line.strip()))
    except (serial.SerialException, OSError) as error:
        serial_error = str(error)
    finally:
        port.close()
        result = dict(elapsed_seconds=round(time.perf_counter()-started, 3), requested_seconds=args.seconds, stopped_by_request=bool(args.stop_file and args.stop_file.exists()), serial_error=serial_error, samples=len(records), reboot_count=reboot_count, last_sample_at=last_sample_at, events=events, records=records)
        write_json(args.out.with_suffix('.summary.json'), result)
    print(json.dumps({k:v for k,v in result.items() if k not in ('records', 'events')}))
    if serial_error:
        sys.exit(1)


if __name__ == '__main__':
    main()