# RoboPark Pi Client
# Run on the Raspberry Pi (or any device) to:
#   1. Enroll with the scheduler (first boot) using --enrollment-token
#   2. Persist device_id + device_token to --state-file (default ./device.json)
#   3. Heartbeat every --heartbeat-interval seconds
#   4. Join a LiveKit room (if `livekit` SDK is installed) to publish mic + camera
#      and subscribe to the agent's data channel for motor commands
#   5. Forward motor commands to the local RoboVisionAI_PI motor server
#
# Quick start:
#   python client.py \
#     --scheduler-url http://100.64.1.5:8080 \
#     --enrollment-token <TOKEN-FROM-SCHEDULER-LOGS> \
#     --name pipi \
#     --tailscale-ip 100.64.1.10 \
#     --lan-ip 192.168.1.159 \
#     --motor-server-url http://192.168.1.159:8001 \
#     --livekit-url ws://100.64.1.5:7880
#
# Subsequent runs (token is in device.json):
#   python client.py --scheduler-url http://100.64.1.5:8080

import argparse
import asyncio
import glob
import json
import logging
import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional

import httpx

from motor_bridge import MotorBridge

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
log = logging.getLogger("robopark-pi")


def get_device_inventory() -> dict:
    """Report real Linux media devices so the remote dashboard can select them."""
    inventory = {
        "video": [{"id": "auto", "name": "Auto detect"}, {"id": "none", "name": "Disable camera"}],
        "audio_input": [{"id": "default", "name": "System default input"}],
        "audio_output": [{"id": "default", "name": "System default output"}],
        "platform": sys.platform,
    }
    for path in sorted(glob.glob("/dev/video*")):
        inventory["video"].append({"id": path, "name": path, "backend": "v4l2"})

    try:
        import pyaudio
        pa = pyaudio.PyAudio()
        for index in range(pa.get_device_count()):
            info = pa.get_device_info_by_index(index)
            name = str(info.get("name", f"Audio device {index}"))
            item = {"id": str(index), "name": name, "host_api": str(info.get("hostApi", ""))}
            if info.get("maxInputChannels", 0) > 0:
                inventory["audio_input"].append(item.copy())
            if info.get("maxOutputChannels", 0) > 0:
                inventory["audio_output"].append(item.copy())
        pa.terminate()
    except Exception as exc:
        # Keep the device visible even when PortAudio is not installed; ALSA
        # names are still useful for an operator diagnosing a Pi remotely.
        log.debug("PyAudio inventory unavailable: %s", exc)
        for command, key, kind in (("arecord", "audio_input", "input"), ("aplay", "audio_output", "output")):
            try:
                output = subprocess.run([command, "-L"], capture_output=True, text=True, timeout=3, check=False).stdout
                for name in (line.strip() for line in output.splitlines()):
                    if name and not name.startswith("(") and name not in {x["id"] for x in inventory[key]}:
                        inventory[key].append({"id": name, "name": name, "backend": "alsa", "kind": kind})
            except Exception:
                pass
    return inventory


class PiClient:
    def __init__(self, args: argparse.Namespace):
        self.scheduler_url = args.scheduler_url.rstrip("/")
        self.state_file = Path(args.state_file)
        self.enrollment_token = args.enrollment_token
        self.name = args.name
        self.tailscale_ip = args.tailscale_ip
        self.lan_ip = args.lan_ip
        self.motor_server_url = args.motor_server_url
        self.livekit_url = args.livekit_url
        self.character_id = args.character_id
        self.heartbeat_interval = args.heartbeat_interval
        self.config_poll_interval = args.config_poll_interval
        self.join_room = args.join_room
        self.room_name = args.room_name
        self.video_device = args.video_device
        self.audio_device = args.audio_device
        self.audio_output_device = args.audio_output_device

        self.device_id: Optional[str] = None
        self.device_token: Optional[str] = None
        self.production_mode: bool = False
        self.motor_bridge = MotorBridge(self.motor_server_url, dry_run=args.dry_run_motors)

        self._stop = asyncio.Event()
        self._lk_task: Optional[asyncio.Task] = None

    # ---- credential persistence ----------------------------------------

    def _load_state(self) -> bool:
        if not self.state_file.exists():
            return False
        try:
            data = json.loads(self.state_file.read_text())
            self.device_id = data.get("device_id")
            self.device_token = data.get("device_token")
            # allow CLI overrides to backfill missing fields from a stale state file
            if not self.tailscale_ip: self.tailscale_ip = data.get("tailscale_ip")
            if not self.lan_ip: self.lan_ip = data.get("lan_ip")
            if not self.motor_server_url: self.motor_server_url = data.get("motor_server_url")
            if not self.livekit_url: self.livekit_url = data.get("livekit_url")
            if not self.character_id: self.character_id = data.get("character_id")
            if not self.name: self.name = data.get("name")
            return bool(self.device_id and self.device_token)
        except Exception as e:
            log.warning(f"Could not load state file {self.state_file}: {e}")
            return False

    def _save_state(self):
        data = {
            "device_id": self.device_id,
            "device_token": self.device_token,
            "name": self.name,
            "tailscale_ip": self.tailscale_ip,
            "lan_ip": self.lan_ip,
            "motor_server_url": self.motor_server_url,
            "livekit_url": self.livekit_url,
            "character_id": self.character_id,
            "saved_at": time.time(),
        }
        self.state_file.parent.mkdir(parents=True, exist_ok=True)
        self.state_file.write_text(json.dumps(data, indent=2))
        try:
            os.chmod(self.state_file, 0o600)
        except Exception:
            pass
        log.info(f"Saved credentials to {self.state_file}")

    # ---- scheduler API --------------------------------------------------

    async def enroll(self):
        body = {
            "enrollment_token": self.enrollment_token,
            "name": self.name,
            "tailscale_ip": self.tailscale_ip,
            "lan_ip": self.lan_ip,
            "motor_server_url": self.motor_server_url,
            "livekit_url": self.livekit_url,
            "character_id": self.character_id,
        }
        async with httpx.AsyncClient(timeout=15.0) as c:
            r = await c.post(f"{self.scheduler_url}/api/devices/enroll", json=body)
            if r.status_code != 200:
                raise RuntimeError(f"Enroll failed ({r.status_code}): {r.text}")
            data = r.json()
        self.device_id = data["device_id"]
        self.device_token = data["device_token"]
        log.info(f"Enrolled as device_id={self.device_id}")
        self._save_state()

    async def heartbeat(self):
        body = {
            "status": "online",
            "ip": self.tailscale_ip or self.lan_ip,
            "device_inventory": get_device_inventory(),
        }
        async with httpx.AsyncClient(timeout=10.0) as c:
            r = await c.post(
                f"{self.scheduler_url}/api/devices/{self.device_id}/heartbeat",
                headers={"Authorization": f"Bearer {self.device_token}"},
                json=body,
            )
            if r.status_code == 401:
                log.error("Device token rejected; re-enrollment required")
                # wipe creds so next loop tries to enroll again
                self.device_token = None
                if self.state_file.exists():
                    self.state_file.unlink()
                return None
            r.raise_for_status()
            return r.json()

    async def fetch_config(self) -> dict:
        async with httpx.AsyncClient(timeout=10.0) as c:
            r = await c.get(
                f"{self.scheduler_url}/api/devices/{self.device_id}/config",
                headers={"Authorization": f"Bearer {self.device_token}"},
            )
            r.raise_for_status()
            return r.json()

    # ---- LiveKit (optional) ---------------------------------------------

    async def livekit_loop(self):
        """Join a LiveKit room, publish mic + cam, subscribe to data channel.
        Skips silently if the `livekit` SDK is not installed."""
        try:
            from livekit_bridge import run_livekit  # local module
        except ImportError:
            log.warning(
                "livekit_bridge module not available (livekit SDK missing). "
                "Heartbeat + motor bridge will still work; install the `livekit` "
                "and `livekit-api` packages on the Pi to enable rooms."
            )
            return

        # Default room: robopark-<device_id>. Allow override via --room-name.
        room = (self.room_name if self.room_name and not self.room_name.startswith("robopark-pi-standby")
                else f"robopark-{self.device_id}")
        # LiveKit URL: prefer the one the scheduler told us about, fall back to CLI.
        lk_url = self.livekit_url or None

        await run_livekit(
            livekit_url=lk_url,
            room_name=room,
            identity=f"pi:{self.device_id}",
            on_motor_command=self.motor_bridge.handle_command,
            stop_event=self._stop,
            production_mode_provider=lambda: self.production_mode,
            scheduler_url=self.scheduler_url,
            device_id=self.device_id,
            device_token=self.device_token,
            video_device=self.video_device,
            audio_device=self.audio_device,
            audio_output_device=self.audio_output_device,
        )

    # ---- main loop ------------------------------------------------------

    async def run(self):
        # 1. credentials
        if not self._load_state():
            if not self.enrollment_token:
                log.error(
                    "No saved credentials and no --enrollment-token. "
                    "First boot requires --enrollment-token from the scheduler."
                )
                return
            await self.enroll()

        log.info(f"Device {self.device_id} starting main loop")
        # 2. start LiveKit worker in background (no-op if SDK missing)
        if self.join_room:
            self._lk_task = asyncio.create_task(self.livekit_loop())
        hb_failures = 0
        try:
            while not self._stop.is_set():
                try:
                    hb = await self.heartbeat()
                    hb_failures = 0
                    if hb:
                        self.production_mode = bool(hb.get("production_mode", False))
                        if self.production_mode and self._lk_task is None and self.join_room:
                            log.info("Production mode ON -> joining LiveKit room")
                            self._lk_task = asyncio.create_task(self.livekit_loop())
                        elif not self.production_mode and self._lk_task is not None:
                            log.info("Production mode OFF -> leaving LiveKit room")
                            self._lk_task.cancel()
                            try:
                                await self._lk_task
                            except (asyncio.CancelledError, Exception):
                                pass
                            self._lk_task = None
                except Exception as e:
                    hb_failures += 1
                    log.warning(f"Heartbeat error ({hb_failures}): {e}")

                # periodic config poll (covers cases where production_mode flips but hb is slow)
                try:
                    if int(time.time()) % int(self.config_poll_interval) < int(self.heartbeat_interval):
                        cfg = await self.fetch_config()
                        self.production_mode = bool(cfg.get("production_mode", False))
                        if cfg.get("character_id"):
                            self.character_id = cfg["character_id"]
                        if cfg.get("video_device") is not None:
                            self.video_device = cfg["video_device"]
                        if cfg.get("audio_device") is not None:
                            self.audio_device = cfg["audio_device"]
                        if cfg.get("audio_output_device") is not None:
                            self.audio_output_device = cfg["audio_output_device"]
                except Exception as e:
                    log.debug(f"Config poll error: {e}")

                try:
                    await asyncio.wait_for(self._stop.wait(), timeout=self.heartbeat_interval)
                except asyncio.TimeoutError:
                    pass
        finally:
            self._stop.set()
            if self._lk_task:
                self._lk_task.cancel()
                try:
                    await self._lk_task
                except Exception:
                    pass


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="RoboPark Pi Client")
    p.add_argument("--scheduler-url", required=True, help="e.g. http://100.64.1.5:8080")
    p.add_argument("--enrollment-token", default=None,
                   help="One-time enrollment token (first boot only)")
    p.add_argument("--state-file", default="./device.json",
                   help="Where to persist device_id + device_token")
    p.add_argument("--name", default=None, help="Device display name")
    p.add_argument("--tailscale-ip", default=None)
    p.add_argument("--lan-ip", default=None)
    p.add_argument("--motor-server-url", default=None,
                   help="Base URL of the local RoboVisionAI_PI motor server")
    p.add_argument("--livekit-url", default=None, help="e.g. ws://100.64.1.5:7880")
    p.add_argument("--character-id", default=None, help="Override character assignment")
    p.add_argument("--room-name", default="robopark-pi-standby",
                   help="LiveKit room to join")
    p.add_argument("--video-device", default="auto", help="Camera path/index, picamera, or auto")
    p.add_argument("--audio-device", default="default", help="Microphone name/index or default")
    p.add_argument("--audio-output-device", default="default", help="Speaker name/index or default")
    p.add_argument("--heartbeat-interval", type=float, default=10.0)
    p.add_argument("--config-poll-interval", type=float, default=30.0)
    p.add_argument("--join-room", action="store_true",
                   help="Attempt to join the LiveKit room (requires livekit SDK)")
    p.add_argument("--dry-run-motors", action="store_true",
                   help="Log motor commands without calling the motor server")
    return p.parse_args()


async def _amain():
    args = parse_args()
    client = PiClient(args)
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            loop.add_signal_handler(sig, client._stop.set)
        except NotImplementedError:
            pass  # Windows

    try:
        await client.run()
    except KeyboardInterrupt:
        pass


def main():
    try:
        asyncio.run(_amain())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()