#!/usr/bin/env python3
"""
ZMQ (JPEG) -> Flask MJPEG Web Viewer
- Subscribes to a ZMQ PUB stream (single-part JPEG frames)
- Serves a browser page + MJPEG endpoint

Usage:
  On the robot:
   conda activate teleimager
   cd teleimager
   python -m teleimager.image_server --rs

  On the  PC:
   python web_server.py
   # then open:
   http://<THIS_MACHINE_IP>:8080
 
Tip:
- If you run this ON the robot, set ROBOT_IP="127.0.0.1"
- If you run this on your laptop, set ROBOT_IP to the robot IP (e.g. 192.168.123.164)
"""

import os
import time
import signal
import threading
from typing import Optional, Tuple

import zmq
import cv2
import numpy as np
from flask import Flask, Response, request

# -----------------------
# Config (env override)
# -----------------------
ROBOT_IP = os.getenv("ROBOT_IP", "10.255.254.86")   # set "127.0.0.1" if running on robot
ZMQ_PORT = int(os.getenv("ZMQ_PORT", "55555"))
HTTP_HOST = os.getenv("HTTP_HOST", "0.0.0.0")
HTTP_PORT = int(os.getenv("HTTP_PORT", "8080"))
JPEG_QUALITY = int(os.getenv("JPEG_QUALITY", "80"))
FRAME_TIMEOUT_MS = int(os.getenv("FRAME_TIMEOUT_MS", "1000"))  # poll timeout for ZMQ
SHOW_FPS = os.getenv("SHOW_FPS", "1") == "1"

app = Flask(__name__)

# -----------------------
# ZMQ Subscriber Worker
# -----------------------
class LatestFrameBuffer:
    """Thread-safe buffer storing only the latest JPEG bytes received."""
    def __init__(self):
        self._lock = threading.Lock()
        self._jpg: Optional[bytes] = None
        self._last_ts = 0.0
        self._fps = 0.0
        self._count = 0
        self._t0 = None

    def update(self, jpg: bytes) -> None:
        now = time.monotonic()
        with self._lock:
            self._jpg = jpg
            self._last_ts = now

            # FPS calc
            if self._t0 is None:
                self._t0 = now
            self._count += 1
            if self._count >= 10:
                dt = now - self._t0
                if dt > 0:
                    self._fps = self._count / dt
                self._t0 = now
                self._count = 0

    def get(self) -> Tuple[Optional[bytes], float, float]:
        with self._lock:
            return self._jpg, self._fps, self._last_ts


class ZMQCameraSubscriber(threading.Thread):
    """Background thread that subscribes to ZMQ and stores the latest frame."""
    def __init__(self, host: str, port: int, buffer: LatestFrameBuffer):
        super().__init__(daemon=True)
        self.host = host
        self.port = port
        self.buffer = buffer
        self._stop = threading.Event()

        self.ctx = zmq.Context.instance()
        self.sub = self.ctx.socket(zmq.SUB)
        self.sub.setsockopt(zmq.RCVHWM, 1)     # keep only latest
        self.sub.setsockopt(zmq.LINGER, 0)
        self.sub.setsockopt_string(zmq.SUBSCRIBE, "")
        self.sub.connect(f"tcp://{self.host}:{self.port}")

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

    def run(self):
        while not self._stop.is_set():
            try:
                events = dict(self.poller.poll(timeout=FRAME_TIMEOUT_MS))
                if self.sub in events:
                    msg = self.sub.recv()  # single-part JPEG
                    if msg:
                        self.buffer.update(msg)
                else:
                    # no frame received in timeout
                    continue
            except Exception:
                time.sleep(0.05)

        try:
            self.sub.close()
        except Exception:
            pass

    def stop(self):
        self._stop.set()


buffer = LatestFrameBuffer()
subscriber = ZMQCameraSubscriber(ROBOT_IP, ZMQ_PORT, buffer)
subscriber.start()

# -----------------------
# Helpers
# -----------------------
def decode_jpeg(jpg: bytes) -> Optional[np.ndarray]:
    try:
        arr = np.frombuffer(jpg, dtype=np.uint8)
        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
        return img
    except Exception:
        return None

def overlay_status(img: np.ndarray, fps: float, last_ts: float) -> np.ndarray:
    if img is None:
        return img
    now = time.monotonic()
    age_ms = (now - last_ts) * 1000.0 if last_ts > 0 else 0.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)
    return img

# -----------------------
# MJPEG streaming
# -----------------------
def mjpeg_generator():
    while True:
        jpg, fps, last_ts = buffer.get()
        if jpg is None:
            time.sleep(0.02)
            continue

        # Optional: decode->overlay->re-encode (adds CPU)
        if SHOW_FPS:
            img = decode_jpeg(jpg)
            if img is None:
                time.sleep(0.01)
                continue
            img = overlay_status(img, fps, last_ts)
            ok, out = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), JPEG_QUALITY])
            if not ok:
                continue
            payload = out.tobytes()
        else:
            payload = jpg  # send original JPEG directly (fastest)

        yield (
            b"--frame\r\n"
            b"Content-Type: image/jpeg\r\n"
            b"Cache-Control: no-cache\r\n\r\n" + payload + b"\r\n"
        )

# -----------------------
# Routes
# -----------------------
@app.route("/")
def index():
    # Allow manual override: /?w=1280
    w = request.args.get("w", "1280")
    return f"""
    <html>
    <head>
      <meta name="viewport" content="width=device-width, initial-scale=1.0"/>
      <title>Robot Head Camera</title>
    </head>
    <body style="margin:0;background:#111;">
      <div style="padding:10px;color:#fff;font-family:Arial;">
        <b>Head Camera</b> | ZMQ: {ROBOT_IP}:{ZMQ_PORT} | MJPEG: {HTTP_HOST}:{HTTP_PORT}
      </div>
      <div style="display:flex;justify-content:center;">
        <img src="/video" style="width:100%;max-width:{w}px;height:auto;border:0;"/>
      </div>
    </body>
    </html>
    """

@app.route("/video")
def video():
    return Response(
        mjpeg_generator(),
        mimetype="multipart/x-mixed-replace; boundary=frame"
    )

@app.route("/health")
def health():
    jpg, fps, last_ts = buffer.get()
    ok = jpg is not None
    age = (time.monotonic() - last_ts) if last_ts else None
    return {
        "ok": ok,
        "robot_ip": ROBOT_IP,
        "zmq_port": ZMQ_PORT,
        "fps": fps,
        "age_sec": age,
    }

# -----------------------
# Clean shutdown
# -----------------------
def _shutdown(*_):
    try:
        subscriber.stop()
    except Exception:
        pass
    time.sleep(0.1)
    raise SystemExit(0)

signal.signal(signal.SIGINT, _shutdown)
signal.signal(signal.SIGTERM, _shutdown)

if __name__ == "__main__":
    print(f"[INFO] Subscribing to: tcp://{ROBOT_IP}:{ZMQ_PORT}")
    print(f"[INFO] Open in browser: http://<this_machine_ip>:{HTTP_PORT}")
    app.run(host=HTTP_HOST, port=HTTP_PORT, threaded=True)
