#!/usr/bin/env python3
"""
fleet_status.py — CLI to display live fleet status.

Usage:
    ros2 run saltybot_fleet fleet_status
    ros2 run saltybot_fleet fleet_status -- --once

Subscribes to:
  /fleet/status           (std_msgs/String JSON)  — coordinator health
  /fleet/robots           (RobotInfo)             — per-robot status
  /fleet/exploration_frontiers (FrontierArray)    — open frontiers
"""

import argparse
import json
import sys

import rclpy
from rclpy.node import Node
from std_msgs.msg import String

from saltybot_fleet_msgs.msg import RobotInfo, FrontierArray


_STATUS_LABELS = {
    RobotInfo.STATUS_IDLE:      'IDLE',
    RobotInfo.STATUS_MAPPING:   'MAPPING',
    RobotInfo.STATUS_EXPLORING: 'EXPLORING',
    RobotInfo.STATUS_RETURNING: 'RETURNING',
    RobotInfo.STATUS_LOST:      'LOST',
}


class FleetStatusCLI(Node):

    def __init__(self, once: bool):
        super().__init__('fleet_status_cli')
        self._once     = once
        self._printed  = False
        self._fleet_status: dict           = {}
        self._robots:       dict           = {}
        self._frontier_count: int          = 0

        self.create_subscription(String,       '/fleet/status',               self._on_status,   10)
        self.create_subscription(RobotInfo,    '/fleet/robots',               self._on_robot,    10)
        self.create_subscription(FrontierArray,'/fleet/exploration_frontiers',self._on_frontiers,10)

        self.create_timer(2.0, self._print_tick)

    def _on_status(self, msg: String) -> None:
        try:
            self._fleet_status = json.loads(msg.data)
        except json.JSONDecodeError:
            pass

    def _on_robot(self, msg: RobotInfo) -> None:
        self._robots[msg.robot_id] = msg

    def _on_frontiers(self, msg: FrontierArray) -> None:
        self._frontier_count = len(msg.frontiers)

    def _print_tick(self) -> None:
        print('\033[H\033[J', end='')   # clear screen
        print('═' * 60)
        print('  SaltyBot Fleet Status')
        print('═' * 60)
        print(f'  Robots in fleet : {len(self._robots)}')
        print(f'  Open frontiers  : {self._frontier_count}')
        print()

        for rid, info in sorted(self._robots.items()):
            status_str = _STATUS_LABELS.get(info.status, '?')
            bat        = f'{info.battery_pct:.0f}%' if info.battery_pct >= 0 else '?'
            frontier   = info.active_frontier or '—'
            print(f'  {rid:<20} {status_str:<12} bat={bat:<6} frontier={frontier}')

        if self._fleet_status:
            print()
            print(f'  Coordinator: {json.dumps(self._fleet_status, indent=2)}')
        print('═' * 60)

        if self._once and not self._printed:
            self._printed = True
            raise SystemExit(0)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--once', action='store_true', help='Print once and exit')
    args, ros_args = parser.parse_known_args()
    rclpy.init(args=ros_args)
    node = FleetStatusCLI(once=args.once)
    try:
        rclpy.spin(node)
    except SystemExit:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
