#!/usr/bin/env python3

import asyncio
import json
import math
import threading
import time
from collections import deque
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

import msgpack
import rclpy
from aiohttp import web
from diagnostic_msgs.msg import DiagnosticArray
from geometry_msgs.msg import Twist
from rclpy.node import Node
from sensor_msgs.msg import Imu, NavSatFix
from std_msgs.msg import Float32, Bool, String


class TelemetryBuffer:
    """Circular buffer for 5-minute history at 2Hz."""

    def __init__(self, duration_seconds=300):
        self.max_samples = int(duration_seconds * 2)
        self.buffer = deque(maxlen=self.max_samples)

    def add(self, data: Dict):
        self.buffer.append({**data, 'ts': time.time()})

    def get_history(self):
        return list(self.buffer)


class RemoteMonitorNode(Node):
    """WebSocket relay for remote monitoring with telemetry aggregation."""

    def __init__(self):
        super().__init__('saltybot_remote_monitor')

        self.declare_parameter('ws_port', 9091)
        self.declare_parameter('auth_token', 'default_token')
        self.declare_parameter('update_rate_hz', 2)
        self.declare_parameter('enable_msgpack', True)

        self.ws_port = self.get_parameter('ws_port').value
        self.auth_token = self.get_parameter('auth_token').value
        self.update_rate_hz = self.get_parameter('update_rate_hz').value
        self.enable_msgpack = self.get_parameter('enable_msgpack').value

        self.battery = {'voltage': 0.0, 'current': 0.0, 'soc': 0.0}
        self.motors = {'left': 0.0, 'right': 0.0}
        self.imu = {'roll': 0.0, 'pitch': 0.0, 'yaw': 0.0}
        self.gps = {'lat': 0.0, 'lon': 0.0, 'alt': 0.0}
        self.health = {'cpu_temp': 0.0, 'gpu_temp': 0.0, 'ram_pct': 0.0, 'disk_pct': 0.0}
        self.social = {'is_speaking': False, 'face_id': None}
        self.alerts = []
        self.history_buffer = TelemetryBuffer()
        self.ws_clients = set()
        self.loop = None

        self.create_subscription(DiagnosticArray, '/diagnostics', self.diag_callback, 10)
        self.create_subscription(Imu, '/saltybot/imu', self.imu_callback, 10)
        self.create_subscription(NavSatFix, '/gps/fix', self.gps_callback, 10)
        self.create_subscription(String, '/saltybot/balance_state', self.balance_callback, 10)
        self.create_subscription(Bool, '/social/speech/is_speaking', self.social_speech_callback, 10)
        self.create_subscription(String, '/social/face/active', self.social_face_callback, 10)

        period = 1.0 / self.update_rate_hz
        self.update_timer = self.create_timer(period, self.telemetry_update_callback)

        self.ws_thread = threading.Thread(target=self.start_ws_server, daemon=True)
        self.ws_thread.start()

        self.get_logger().info(
            f'Remote monitor initialized on port {self.ws_port}, '
            f'update rate {self.update_rate_hz}Hz, msgpack={self.enable_msgpack}'
        )

    def diag_callback(self, msg: DiagnosticArray):
        """Parse diagnostics for battery and system health."""
        for status in msg.status:
            kv = {pair.key: pair.value for pair in status.values}

            if kv.get('battery_voltage_v'):
                self.battery = {
                    'voltage': float(kv.get('battery_voltage_v', 0)),
                    'current': float(kv.get('battery_current_a', 0)),
                    'soc': float(kv.get('battery_soc_pct', 0)),
                }
                if self.battery['soc'] < 15:
                    self.add_alert('CRITICAL', 'Low battery', self.battery['soc'])

            if kv.get('cpu_temp_c'):
                self.health = {
                    'cpu_temp': float(kv.get('cpu_temp_c', 0)),
                    'gpu_temp': float(kv.get('gpu_temp_c', 0)),
                    'ram_pct': float(kv.get('ram_pct', 0)),
                    'disk_pct': float(kv.get('disk_pct', 0)),
                }

    def imu_callback(self, msg: Imu):
        """Extract roll/pitch/yaw from quaternion."""
        q = msg.orientation
        self.imu = self.quat_to_euler(q.x, q.y, q.z, q.w)

    def gps_callback(self, msg: NavSatFix):
        """Store GPS data."""
        self.gps = {
            'lat': msg.latitude,
            'lon': msg.longitude,
            'alt': msg.altitude,
        }

    def balance_callback(self, msg: String):
        """Parse balance state for motor commands."""
        try:
            state = json.loads(msg.data)
            cmd = state.get('motor_cmd', 0)
            norm = max(-1, min(1, cmd / 1000))
            self.motors = {'left': norm, 'right': norm}

            if abs(norm) > 0.9:
                self.add_alert('WARNING', 'High motor load', norm)
        except Exception:
            pass

    def social_speech_callback(self, msg: Bool):
        """Update social speaking state."""
        self.social['is_speaking'] = msg.data

    def social_face_callback(self, msg: String):
        """Update recognized face ID."""
        self.social['face_id'] = msg.data

    @staticmethod
    def quat_to_euler(qx, qy, qz, qw):
        """Convert quaternion to Euler angles."""
        sinr_cosp = 2 * (qw * qx + qy * qz)
        cosr_cosp = 1 - 2 * (qx * qx + qy * qy)
        roll = math.atan2(sinr_cosp, cosr_cosp)

        sinp = 2 * (qw * qy - qz * qx)
        pitch = math.asin(max(-1, min(1, sinp)))

        siny_cosp = 2 * (qw * qz + qx * qy)
        cosy_cosp = 1 - 2 * (qy * qy + qz * qz)
        yaw = math.atan2(siny_cosp, cosy_cosp)

        return {
            'roll': (roll * 180) / math.pi,
            'pitch': (pitch * 180) / math.pi,
            'yaw': (yaw * 180) / math.pi,
        }

    def add_alert(self, level: str, message: str, value: Optional[float] = None):
        """Add alert with timestamp."""
        alert = {
            'ts': time.time(),
            'level': level,
            'message': message,
            'value': value,
        }
        self.alerts.append(alert)
        if len(self.alerts) > 50:
            self.alerts.pop(0)

    def telemetry_update_callback(self):
        """Periodic telemetry aggregation and broadcast."""
        telemetry = {
            'timestamp': datetime.now().isoformat(),
            'battery': self.battery,
            'motors': self.motors,
            'imu': self.imu,
            'gps': self.gps,
            'health': self.health,
            'social': self.social,
            'alerts': self.alerts[-10:],
        }

        self.history_buffer.add(telemetry)

        if self.ws_clients and self.loop:
            try:
                asyncio.run_coroutine_threadsafe(
                    self.broadcast_telemetry(telemetry),
                    self.loop
                )
            except Exception as e:
                self.get_logger().warn(f'Broadcast error: {e}')

    async def broadcast_telemetry(self, telemetry: Dict):
        """Send telemetry to all connected WebSocket clients."""
        if self.enable_msgpack:
            data = msgpack.packb(telemetry)
        else:
            data = json.dumps(telemetry).encode('utf-8')

        disconnected = set()
        for ws in self.ws_clients:
            try:
                await ws.send_bytes(data)
            except Exception:
                disconnected.add(ws)

        self.ws_clients -= disconnected

    async def websocket_handler(self, request):
        """Handle WebSocket connections."""
        token = request.rel_url.query.get('token')

        if token != self.auth_token:
            return web.Response(status=401, text='Unauthorized')

        ws = web.WebSocketResponse()
        await ws.prepare(request)
        self.ws_clients.add(ws)

        try:
            history = self.history_buffer.get_history()
            if self.enable_msgpack:
                await ws.send_bytes(msgpack.packb({'type': 'history', 'data': history}))
            else:
                await ws.send_str(json.dumps({'type': 'history', 'data': history}))

            async for msg in ws:
                if msg.type == web.WSMsgType.TEXT:
                    try:
                        req = json.loads(msg.data)
                        if req.get('cmd') == 'get_history':
                            history = self.history_buffer.get_history()
                            if self.enable_msgpack:
                                await ws.send_bytes(msgpack.packb({'type': 'history', 'data': history}))
                            else:
                                await ws.send_str(json.dumps({'type': 'history', 'data': history}))
                    except Exception as e:
                        self.get_logger().error(f'WebSocket error: {e}')

        except Exception as e:
            self.get_logger().warn(f'WebSocket connection lost: {e}')
        finally:
            self.ws_clients.discard(ws)

        return ws

    async def index_handler(self, request):
        """Serve mobile-friendly HTML UI."""
        html = Path(__file__).parent.parent / 'static' / 'index.html'
        if html.exists():
            return web.FileResponse(html)
        return web.Response(text='UI not found', status=404)

    def start_ws_server(self):
        """Start aiohttp WebSocket server."""
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        app = web.Application()
        app.router.add_get('/', self.index_handler)
        app.router.add_get('/ws', self.websocket_handler)

        runner = web.AppRunner(app)
        self.loop.run_until_complete(runner.setup())
        site = web.TCPSite(runner, '0.0.0.0', self.ws_port)
        self.loop.run_until_complete(site.start())

        self.get_logger().info(f'WebSocket server running on port {self.ws_port}')
        self.loop.run_forever()


def main(args=None):
    rclpy.init(args=args)
    node = RemoteMonitorNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
