#!/usr/bin/env python3
"""
Pulse oximeter UDP listener.
Receives heart-rate (BPM) and SpO2 readings from the device and prints them.
Run:  python3 receive.py
"""

import socket
import datetime
import json
import csv
import sys
import os

HOST = "0.0.0.0"
PORT = 5005

def parse_payload(raw: bytes) -> dict | None:
    """Try JSON first, then common key=value or CSV formats."""
    text = raw.decode(errors="replace").strip()
    if not text:
        return None

    # JSON
    if text.startswith("{"):
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass

    # key=value pairs (e.g. BPM=72,SPO2=98)
    if "=" in text:
        try:
            return dict(part.split("=", 1) for part in text.replace(";", ",").split(",") if "=" in part)
        except Exception:
            pass

    # Bare CSV: treat first field as BPM, second as SpO2
    parts = text.split(",")
    if len(parts) >= 2 and all(p.strip().lstrip("-").isdigit() for p in parts[:2]):
        return {"BPM": parts[0].strip(), "SpO2": parts[1].strip()}

    # Unknown — return raw
    return {"raw": text}


def main():
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((HOST, PORT))

    log_path = os.path.join(os.path.dirname(__file__), "readings.csv")
    write_header = not os.path.exists(log_path)
    log_file = open(log_path, "a", newline="")
    writer = csv.writer(log_file)
    if write_header:
        writer.writerow(["timestamp", "bpm", "spo2", "raw"])

    print(f"Listening on UDP {HOST}:{PORT}  |  Logging to {log_path}")
    print(f"{'TIME':12}  {'BPM':>6}  {'SpO2 %':>7}  STATUS")
    print("-" * 52)

    try:
        while True:
            data, addr = sock.recvfrom(4096)
            ts = datetime.datetime.now()
            parsed = parse_payload(data)

            def _get(d, *keys):
                for k in keys:
                    if k in d:
                        return d[k]
                return None

            bpm_raw  = _get(parsed, "bpm", "BPM", "heartrate")
            spo2_raw = _get(parsed, "spo2", "SPO2", "SpO2", "oxygen")
            raw      = parsed.get("raw", "")

            bpm_val  = int(bpm_raw)  if bpm_raw  is not None else None
            spo2_val = int(spo2_raw) if spo2_raw is not None else None

            no_finger = (bpm_val is not None and bpm_val < 0) or \
                        (spo2_val is not None and spo2_val < 0)

            bpm_str  = "---" if (bpm_val  is None or bpm_val  < 0) else str(bpm_val)
            spo2_str = "---" if (spo2_val is None or spo2_val < 0) else str(spo2_val)
            status   = "No finger" if no_finger else (f"({raw})" if raw else "OK")

            print(f"{ts.strftime('%H:%M:%S.%f')[:12]}  {bpm_str:>6}  {spo2_str:>7}  {status}")
            sys.stdout.flush()

            # Log only rows where valid readings exist
            log_bpm  = bpm_val  if (bpm_val  is not None and bpm_val  >= 0) else ""
            log_spo2 = spo2_val if (spo2_val is not None and spo2_val >= 0) else ""
            writer.writerow([ts.isoformat(), log_bpm, log_spo2, data.decode(errors="replace").strip()])
            log_file.flush()

    except KeyboardInterrupt:
        print("\nStopped.")
    finally:
        sock.close()
        log_file.close()


if __name__ == "__main__":
    main()