#!/usr/bin/env python3
"""
Laptop ZMQ JPEG viewer (OpenCV window)
- Subscribes to robot ZMQ PUB stream (single-part JPEG)
- Displays live video in a popup window
- Optional FPS + frame age overlay

Run:
  pip install pyzmq opencv-python numpy
  python laptop_zmq_viewer.py --robot 192.168.123.164 --port 55555

Keys:
  q  -> quit
  f  -> toggle overlay
"""

import time
import argparse
import zmq
import cv2
import numpy as np


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--robot", default="10.255.254.86", help="Robot IP that publishes ZMQ JPEG")
    ap.add_argument("--port", type=int, default=55555, help="ZMQ port (from cam_config_client.yaml)")
    ap.add_argument("--hwm", type=int, default=1, help="ZMQ high-water-mark (keep latest frames)")
    ap.add_argument("--timeout_ms", type=int, default=1000, help="ZMQ poll timeout")
    ap.add_argument("--window", default="G1 Head Camera", help="OpenCV window name")
    args = ap.parse_args()

    robot_ip = args.robot
    port = args.port

    # -----------------------
    # ZMQ subscriber
    # -----------------------
    ctx = zmq.Context.instance()
    sub = ctx.socket(zmq.SUB)
    sub.setsockopt(zmq.RCVHWM, args.hwm)
    sub.setsockopt(zmq.LINGER, 0)
    sub.setsockopt_string(zmq.SUBSCRIBE, "")
    sub.connect(f"tcp://{robot_ip}:{port}")

    poller = zmq.Poller()
    poller.register(sub, zmq.POLLIN)

    print(f"[INFO] Subscribing to tcp://{robot_ip}:{port}")
    print("[INFO] Press 'q' to quit, 'f' to toggle overlay")

    # -----------------------
    # Stats
    # -----------------------
    show_overlay = True
    last_frame_ts = 0.0
    fps = 0.0
    fps_count = 0
    t0 = time.monotonic()

    cv2.namedWindow(args.window, cv2.WINDOW_NORMAL)

    while True:
        events = dict(poller.poll(timeout=args.timeout_ms))
        if sub not in events:
            # no frame received
            # show a black frame with warning (optional)
            black = np.zeros((480, 640, 3), dtype=np.uint8)
            cv2.putText(black, "No frames...", (20, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2, cv2.LINE_AA)
            cv2.imshow(args.window, black)
        else:
            jpg = sub.recv()  # single-part JPEG bytes
            arr = np.frombuffer(jpg, dtype=np.uint8)
            img = cv2.imdecode(arr, cv2.IMREAD_COLOR)

            if img is None:
                continue

            now = time.monotonic()
            last_frame_ts = now

            # FPS estimate (every 10 frames)
            fps_count += 1
            if fps_count >= 10:
                dt = now - t0
                if dt > 0:
                    fps = fps_count / dt
                t0 = now
                fps_count = 0

            if show_overlay:
                age_ms = (time.monotonic() - last_frame_ts) * 1000.0
                text = f"FPS: {fps:.1f} | Age: {age_ms:.0f} ms"
                cv2.putText(img, text, (12, 28),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)

            cv2.imshow(args.window, img)

        key = cv2.waitKey(1) & 0xFF
        if key == ord("q"):
            break
        if key == ord("f"):
            show_overlay = not show_overlay

    cv2.destroyAllWindows()
    sub.close()
    ctx.term()


if __name__ == "__main__":
    main()