#!/usr/bin/env python3
"""
HoboStreamer × Cozmo — Let viewers drive your robot live
==========================================================
Connects a Cozmo robot (via PyCozmo) to HoboStreamer's control system so
viewers can press dashboard buttons to drive, animate, and interact with
the robot in real time. Optionally streams Cozmo's front camera to your
live stream via ffmpeg + jsmpeg.

How it works
------------
  1. This script connects PyCozmo to the Cozmo robot over Wi-Fi.
  2. It connects to HoboStreamer's control WebSocket as a hardware client
     using your stream key.
  3. Viewers click control buttons in the live stream → HoboStreamer
     routes those commands here → PyCozmo executes them on the robot.
  4. Optionally, Cozmo's camera feed is piped through ffmpeg to the
     jsmpeg ingest endpoint so viewers can see what Cozmo sees.

Setup
-----
  1. Install PyCozmo:
       pip install pycozmo pillow
       pycozmo_resources.py download

  2. Install websocket-client:
       pip install websocket-client

  3. Connect to Cozmo's Wi-Fi:
       - Place Cozmo on charger to wake it.
       - Raise and lower the lift — Cozmo displays its Wi-Fi password.
       - Join Cozmo's Wi-Fi network on your PC/Mac.

  4. Get your HoboStreamer stream key:
       hobostreamer.com → Dashboard → Stream Key section

  5. Go live on HoboStreamer (Dashboard → Go Live, protocol: jsmpeg or rtmp).
     Note the Stream ID in the Active Streams list.

  6. Add control buttons via the dashboard or via hobostreamer_example.py:
       Recommended commands: forward, backward, left, right, stop,
                             wave, happy, sad, look_up, look_down, say

  7. Run this script:
       python hobostreamer_cozmo.py --key YOUR_STREAM_KEY

  8. Optional — stream Cozmo's camera via jsmpeg (requires ffmpeg):
       python hobostreamer_cozmo.py --key YOUR_STREAM_KEY --stream-camera

Usage
-----
  python hobostreamer_cozmo.py --key YOUR_STREAM_KEY [options]

Options
-------
  --key KEY            HoboStreamer stream key (required)
  --host HOST          HoboStreamer server (default: hobostreamer.com)
  --port PORT          Server port (default: 443)
  --no-ssl             Disable SSL (for local dev server)
  --stream-camera      Pipe Cozmo's camera to ffmpeg for jsmpeg streaming
  --camera-fps N       Camera capture FPS (default: 15)
  --drive-speed N      Wheel speed 0–300 mm/s (default: 150)
  --drive-duration N   Seconds to drive per command (default: 0.5)
  --turn-duration N    Seconds to turn per command (default: 0.4)
  --sim                Skip PyCozmo connection (simulation mode for testing)
"""

import argparse
import json
import subprocess
import sys
import signal
import time
import threading
import io

try:
    import websocket
except ImportError:
    print("[ERROR] websocket-client not installed. Run: pip install websocket-client")
    sys.exit(1)

PYCOZMO_AVAILABLE = False
try:
    import pycozmo
    from pycozmo.lights import Color
    PYCOZMO_AVAILABLE = True
except ImportError:
    pass

DEFAULT_HOST = "hobostreamer.com"
DEFAULT_PORT = 443
DRIVE_SPEED    = 150.0   # mm/s
DRIVE_DURATION = 0.5     # seconds
TURN_DURATION  = 0.4     # seconds


# ═══════════════════════════════════════════════════════════════
# Cozmo Controller
# ═══════════════════════════════════════════════════════════════

class CozmoController:
    """
    Wraps a PyCozmo client and maps HoboStreamer commands to robot actions.

    Command reference
    -----------------
      forward      — drive forward
      backward     — drive backward
      left         — turn left (spin in place)
      right        — turn right (spin in place)
      stop         — stop all motors
      wave         — play wave animation
      happy        — happy reaction animation + backpack lights green
      sad          — sad reaction animation + backpack lights blue
      look_up      — tilt head up
      look_down    — tilt head down
      look_center  — center head tilt
      lift_up      — raise lift
      lift_down    — lower lift
      lift_center  — center lift
      say <text>   — speak text (limited to 80 chars)
      lights_on    — backpack LEDs white
      lights_off   — backpack LEDs off
    """

    def __init__(self, cli=None, drive_speed=DRIVE_SPEED,
                 drive_duration=DRIVE_DURATION, turn_duration=TURN_DURATION,
                 sim=False):
        self.cli = cli
        self.sim = sim
        self.drive_speed   = float(drive_speed)
        self.drive_duration = float(drive_duration)
        self.turn_duration  = float(turn_duration)
        self._lock = threading.Lock()

    # ── Internal helpers ────────────────────────────────────────

    def _drive(self, l_speed, r_speed, duration):
        if self.sim or self.cli is None:
            print(f"    [SIM] drive l={l_speed} r={r_speed} for {duration}s")
            return
        with self._lock:
            self.cli.drive_wheels(
                lwheel_speed=l_speed,
                rwheel_speed=r_speed,
                duration=duration,
            )

    def _set_head(self, angle):
        """angle in radians: ~-0.44 (down) to ~0.78 (up)"""
        if self.sim or self.cli is None:
            print(f"    [SIM] head angle={angle:.2f}")
            return
        with self._lock:
            self.cli.set_head_angle(angle=angle)

    def _set_lift(self, height):
        """height 0.0 (down) to 1.0 (up)"""
        if self.sim or self.cli is None:
            print(f"    [SIM] lift height={height:.2f}")
            return
        with self._lock:
            self.cli.set_lift_height(height=height)

    def _play_animation(self, name):
        if self.sim or self.cli is None:
            print(f"    [SIM] animation: {name}")
            return
        try:
            with self._lock:
                self.cli.play_anim(name=name)
        except Exception as e:
            print(f"    [ANIM] Failed to play {name}: {e}")

    def _set_backpack_lights(self, r=0, g=0, b=0, on=True):
        if self.sim or self.cli is None:
            state = f"r={r} g={g} b={b}" if on else "off"
            print(f"    [SIM] backpack lights {state}")
            return
        try:
            if on:
                color = pycozmo.lights.Color(r=r, g=g, b=b)
                self.cli.set_all_backpack_lights(pycozmo.lights.Light(on_color=color))
            else:
                self.cli.set_all_backpack_lights(pycozmo.lights.Light())
        except Exception:
            pass

    # ── Command dispatch ────────────────────────────────────────

    def handle_command(self, command, data, from_user):
        """
        Called by the WebSocket listener for every viewer button press.

        Parameters
        ----------
        command   : str  — command name defined in the dashboard button
        data      : dict — optional extra payload (may be {})
        from_user : str  — viewer username or 'anonymous'
        """
        print(f"  [CMD] {command!r} from {from_user}")
        cmd = command.strip().lower()

        if cmd == "forward":
            self._drive(self.drive_speed, self.drive_speed, self.drive_duration)

        elif cmd == "backward":
            self._drive(-self.drive_speed, -self.drive_speed, self.drive_duration)

        elif cmd == "left":
            self._drive(-self.drive_speed, self.drive_speed, self.turn_duration)

        elif cmd == "right":
            self._drive(self.drive_speed, -self.drive_speed, self.turn_duration)

        elif cmd == "stop":
            self._drive(0, 0, 0)

        elif cmd == "wave":
            self._play_animation("anim_greeting_wave_01")

        elif cmd == "happy":
            self._play_animation("anim_pounce_happy_02")
            self._set_backpack_lights(r=0, g=255, b=0)

        elif cmd == "sad":
            self._play_animation("anim_bored_event_01")
            self._set_backpack_lights(r=0, g=0, b=255)

        elif cmd == "look_up":
            self._set_head(0.6)

        elif cmd == "look_down":
            self._set_head(-0.2)

        elif cmd == "look_center":
            self._set_head(0.0)

        elif cmd == "lift_up":
            self._set_lift(1.0)

        elif cmd == "lift_down":
            self._set_lift(0.0)

        elif cmd == "lift_center":
            self._set_lift(0.5)

        elif cmd.startswith("say ") or cmd == "say":
            # Command can be 'say' with text in data, or 'say Hello there' inline
            if cmd.startswith("say "):
                text = command[4:].strip()[:80]
            else:
                text = str(data.get("text", "Hello!"))[:80]
            print(f"    → say: {text!r}")
            if not self.sim and self.cli is not None:
                try:
                    self.cli.say_text(text)
                except Exception as e:
                    print(f"    [TTS ERROR] {e}")

        elif cmd == "lights_on":
            self._set_backpack_lights(r=255, g=255, b=255, on=True)

        elif cmd == "lights_off":
            self._set_backpack_lights(on=False)

        else:
            print(f"    → unknown command: {command!r}")


# ═══════════════════════════════════════════════════════════════
# Camera Streaming — pipe Cozmo's camera to ffmpeg (jsmpeg)
# ═══════════════════════════════════════════════════════════════

def start_camera_stream(cli, stream_key, host, port, ssl, fps=15):
    """
    Capture frames from Cozmo's front camera and pipe them into ffmpeg,
    which encodes them as MPEG-1 and sends to HoboStreamer's jsmpeg relay.

    The stream will appear as Cozmo's POV in the live stream.

    Requirements: ffmpeg installed and in PATH.
    """
    try:
        import PIL.Image
    except ImportError:
        print("[CAM] Pillow not installed — camera stream unavailable. Run: pip install pillow")
        return

    scheme_tcp = "tcp"
    port_str   = f":{port}" if port not in (80, 443) else (":80" if not ssl else ":443")
    tcp_url    = f"{scheme_tcp}://{host}{port_str}/jsmpeg/{stream_key}"

    # Cozmo camera is 320×240
    cam_w, cam_h = 320, 240

    ffmpeg_cmd = [
        "ffmpeg", "-y",
        "-f", "rawvideo",
        "-pixel_format", "rgb24",
        "-video_size", f"{cam_w}x{cam_h}",
        "-framerate", str(fps),
        "-i", "pipe:0",
        "-an",
        "-vcodec", "mpeg1video",
        "-b:v", "400k",
        "-r", str(fps),
        "-f", "mpegts",
        tcp_url,
    ]

    print(f"[CAM] Starting Cozmo camera stream → {tcp_url}")
    print(f"[CAM] ffmpeg: {' '.join(ffmpeg_cmd)}\n")

    try:
        proc = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE)
    except FileNotFoundError:
        print("[CAM] ffmpeg not found. Install ffmpeg to enable camera streaming.")
        return

    def on_camera_image(cli, image):
        """Called by PyCozmo for each new camera frame."""
        try:
            # image is a PIL.Image in RGB
            frame_bytes = image.tobytes()
            proc.stdin.write(frame_bytes)
        except (BrokenPipeError, OSError):
            pass

    if cli is not None:
        cli.add_handler(pycozmo.event.EvtNewRawCameraImage, on_camera_image)
        cli.enable_camera(enable=True, color=True)
        print("[CAM] Camera enabled")
    else:
        print("[CAM] Simulation mode — no real camera frames")

    return proc


# ═══════════════════════════════════════════════════════════════
# HoboStreamer WebSocket Hardware Client
# ═══════════════════════════════════════════════════════════════

def listen_for_commands(stream_key, controller, host, port, ssl):
    """Connect to HoboStreamer as a hardware client and dispatch to controller."""
    scheme = "wss" if ssl else "ws"
    port_str = f":{port}" if (ssl and port != 443) or (not ssl and port != 80) else ""
    url = f"{scheme}://{host}{port_str}/ws/control?mode=hardware&stream_key={stream_key}"
    reconnect_delay = 5
    running = True

    def on_open(ws):
        print(f"[WS] Connected to HoboStreamer control server")
        print(f"[WS] Waiting for viewer commands...\n")

    def on_message(ws, raw):
        try:
            msg = json.loads(raw)
            msg_type = msg.get("type")

            if msg_type == "connected":
                print("[WS] Hardware client registered — you're live!")

            elif msg_type == "command":
                controller.handle_command(
                    command=msg.get("command", ""),
                    data=msg.get("data") or {},
                    from_user=msg.get("from_user", "anonymous"),
                )

            elif msg_type == "error":
                print(f"[WS ERROR] {msg.get('message')}")

            elif msg_type == "ping":
                ws.send(json.dumps({"type": "pong"}))

        except json.JSONDecodeError:
            print(f"[WS] Malformed message: {raw[:200]}")

    def on_error(ws, error):
        print(f"[WS ERROR] {error}")

    def on_close(ws, code, reason):
        print(f"[WS] Disconnected (code={code}, reason={reason})")
        if running:
            print(f"[WS] Reconnecting in {reconnect_delay}s...")
            time.sleep(reconnect_delay)
            connect()

    def connect():
        ws = websocket.WebSocketApp(
            url,
            on_open=on_open,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close,
        )
        ws.run_forever()

    def stop(sig, frame):
        nonlocal running
        running = False
        print("\n[WS] Shutting down.")
        sys.exit(0)

    signal.signal(signal.SIGINT, stop)
    signal.signal(signal.SIGTERM, stop)
    connect()


# ═══════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="HoboStreamer × Cozmo — Let viewers drive your robot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--key",             required=True, help="Your HoboStreamer stream key")
    parser.add_argument("--host",            default=DEFAULT_HOST)
    parser.add_argument("--port",            type=int, default=DEFAULT_PORT)
    parser.add_argument("--no-ssl",          action="store_true", help="Disable SSL (local dev)")
    parser.add_argument("--stream-camera",   action="store_true",
                        help="Stream Cozmo's camera to the jsmpeg endpoint")
    parser.add_argument("--camera-fps",      type=int, default=15, help="Camera FPS (default 15)")
    parser.add_argument("--drive-speed",     type=float, default=DRIVE_SPEED,
                        help=f"Drive speed mm/s (default {DRIVE_SPEED})")
    parser.add_argument("--drive-duration",  type=float, default=DRIVE_DURATION,
                        help=f"Drive duration seconds (default {DRIVE_DURATION})")
    parser.add_argument("--turn-duration",   type=float, default=TURN_DURATION,
                        help=f"Turn duration seconds (default {TURN_DURATION})")
    parser.add_argument("--sim",             action="store_true",
                        help="Simulation mode — skip PyCozmo connection (for testing without a robot)")
    args = parser.parse_args()

    ssl = not args.no_ssl

    print("""
 ██████╗ ██████╗ ███████╗███╗   ███╗ ██████╗
██╔════╝██╔═══██╗╚══███╔╝████╗ ████║██╔═══██╗
██║     ██║   ██║  ███╔╝ ██╔████╔██║██║   ██║
██║     ██║   ██║ ███╔╝  ██║╚██╔╝██║██║   ██║
╚██████╗╚██████╔╝███████╗██║ ╚═╝ ██║╚██████╔╝
 ╚═════╝ ╚═════╝ ╚══════╝╚═╝     ╚═╝ ╚═════╝
   HoboStreamer × Cozmo Robot Controller
""")

    if args.sim:
        print("[INFO] Simulation mode — Cozmo not connected")

    # ── Connect to Cozmo ────────────────────────────────────────
    cli = None
    if not args.sim:
        if not PYCOZMO_AVAILABLE:
            print("[ERROR] PyCozmo not installed. Run: pip install pycozmo")
            print("[INFO]  To run without a robot, use: --sim")
            sys.exit(1)

        print("[COZMO] Connecting to Cozmo over Wi-Fi...")
        print("[COZMO] Make sure you're connected to Cozmo's Wi-Fi network first.")
        print("[COZMO]   (Place on charger → raise/lower lift to see Wi-Fi password)\n")

        try:
            cli = pycozmo.Client()
            cli.start()
            cli.connect()
            cli.wait_for_robot()
            print("[COZMO] Connected!\n")
        except Exception as e:
            print(f"[ERROR] Failed to connect to Cozmo: {e}")
            print("[INFO]  Use --sim to run without a robot.")
            sys.exit(1)
    else:
        print("[COZMO] Skipping robot connection (--sim)\n")

    # ── Build controller ────────────────────────────────────────
    controller = CozmoController(
        cli=cli,
        drive_speed=args.drive_speed,
        drive_duration=args.drive_duration,
        turn_duration=args.turn_duration,
        sim=args.sim,
    )

    # ── Start camera stream thread (optional) ──────────────────
    cam_proc = None
    if args.stream_camera:
        cam_proc = start_camera_stream(
            cli=cli,
            stream_key=args.key,
            host=args.host,
            port=args.port,
            ssl=ssl,
            fps=args.camera_fps,
        )

    print(f"[INFO] Connecting to HoboStreamer: {args.host}:{args.port}")
    print(f"[INFO] Stream key: {args.key[:6]}{'*' * (len(args.key) - 6)}\n")
    print("[INFO] Supported commands: forward, backward, left, right, stop,")
    print("[INFO]   wave, happy, sad, look_up, look_down, look_center,")
    print("[INFO]   lift_up, lift_down, lift_center, say <text>,")
    print("[INFO]   lights_on, lights_off\n")

    # ── Connect to HoboStreamer & block ─────────────────────────
    try:
        listen_for_commands(
            stream_key=args.key,
            controller=controller,
            host=args.host,
            port=args.port,
            ssl=ssl,
        )
    finally:
        if cam_proc:
            cam_proc.terminate()
        if cli:
            try:
                cli.disconnect()
                cli.stop()
            except Exception:
                pass
        print("[INFO] Goodbye.")


if __name__ == "__main__":
    main()