#!/usr/bin/env python3
"""
ESP32-S3 LD2453 Radar Visualizer for Orange Pi
===============================================
- Kết nối WebSocket tới ESP32 để nhận dữ liệu radar realtime
- Vẽ zone (3 vùng hình chữ nhật) và vị trí các target (up to 3)
- Hiển thị tọa độ, tốc độ, khoảng cách của từng target
- Hỗ trợ chạy headless (chỉ log) hoặc GUI (tkinter/pygame)

Sử dụng:
  python orange_pi_visualizer.py <ESP32_IP>              # GUI mode (tkinter)
  python orange_pi_visualizer.py <ESP32_IP> --headless   # Chỉ in console
  python orange_pi_visualizer.py <ESP32_IP> --mode pygame # Dùng pygame (đồ họa đẹp hơn)
"""

import argparse
import asyncio
import json
import sys
import time
from collections import deque
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

# ============================================================
# WebSocket client
# ============================================================
try:
    import websockets
except ImportError:
    print("Loi: Can install websockets: pip install websockets")
    sys.exit(1)


@dataclass
class Target:
    slot: int
    valid: bool = False
    x_mm: int = 0
    y_mm: int = 0
    speed_cm_s: int = 0
    distance_mm: float = 0.0
    angle_deg: float = 0.0
    last_seen: float = 0.0


@dataclass
class Zone:
    id: int
    x1: int = 0
    y1: int = 0
    x2: int = 0
    y2: int = 0
    occupied: bool = False


@dataclass
class RadarState:
    node_id: str = "RADAR_A"
    sequence: int = 0
    timestamp_ms: int = 0
    target_count: int = 0
    targets: List[Target] = field(default_factory=lambda: [Target(i + 1) for i in range(3)])
    zones: List[Zone] = field(default_factory=lambda: [Zone(i) for i in range(1, 4)])
    frames_received: int = 0
    last_update: float = 0.0

    def update_from_frame(self, frame: dict):
        self.node_id = frame.get("node_id", self.node_id)
        self.sequence = frame.get("sequence", 0)
        self.timestamp_ms = frame.get("timestamp", 0)
        self.target_count = frame.get("target_count", 0)
        now = time.time()

        targets_data = frame.get("targets", [])
        for t_data in targets_data:
            slot = t_data.get("id", 0)
            if 1 <= slot <= 3:
                t = self.targets[slot - 1]
                t.slot = slot
                t.valid = t_data.get("valid", False)
                t.x_mm = t_data.get("x", 0)
                t.y_mm = t_data.get("y", 0)
                t.speed_cm_s = t_data.get("speed", 0)
                t.distance_mm = float(t_data.get("distance", 0))
                t.angle_deg = float(t_data.get("angle", 0))
                if t.valid:
                    t.last_seen = now

        self.frames_received += 1
        self.last_update = now


# ============================================================
# GUI 1: Tkinter (mặc định, có sẵn trên hầu hết hệ thống)
# ============================================================
class TkinterVisualizer:
    # Canvas size (pixel)
    CANVAS_W = 900
    CANVAS_H = 700

    # Radar coordinate range (mm) — map vào canvas
    ROOM_X_MIN = -5000  # 5m bên trái radar
    ROOM_X_MAX = 5000   # 5m bên phải
    ROOM_Y_MIN = 0      # radar
    ROOM_Y_MAX = 7000   # 7m hướng về phía trước

    # Colors
    BG_COLOR = "#1e1e2e"
    GRID_COLOR = "#313244"
    AXIS_COLOR = "#89b4fa"
    ZONE_COLORS = ["#a6e3a1", "#fab387", "#f38ba8", "#89dceb"]
    ZONE_OCCUPIED_COLORS = ["#40a02b", "#fe640b", "#d20f39", "#04a5e5"]
    TARGET_COLORS = ["#f9e2af", "#cba6f7", "#74c7ec"]

    def __init__(self, state: RadarState):
        self.state = state
        try:
            import tkinter as tk
        except ImportError:
            print("Loi: tkinter khong co san. Hay cai: sudo apt install python3-tk")
            sys.exit(1)
        self.tk = tk
        self.root = tk.Tk()
        self.root.title(f"ESP32 LD2453 Radar — {state.node_id}")
        self.root.configure(bg=self.BG_COLOR)
        self.root.protocol("WM_DELETE_WINDOW", self._on_close)
        self._running = True

        self._build_ui()

    def _build_ui(self):
        # Header label
        self.header_var = self.tk.StringVar(value="Dang cho du lieu...")
        lbl = self.tk.Label(
            self.root, textvariable=self.header_var,
            font=("Consolas", 12, "bold"),
            fg="#cdd6f4", bg=self.BG_COLOR, pady=8
        )
        lbl.pack(fill="x")

        # Canvas
        self.canvas = self.tk.Canvas(
            self.root, width=self.CANVAS_W, height=self.CANVAS_H,
            bg=self.BG_COLOR, highlightthickness=0
        )
        self.canvas.pack(padx=10, pady=5)

        # Status bar
        self.status_var = self.tk.StringVar(value="Frames: 0")
        self.tk.Label(
            self.root, textvariable=self.status_var,
            font=("Consolas", 10),
            fg="#a6adc8", bg=self.BG_COLOR, pady=6
        ).pack(fill="x")

    def _mm_to_pixel(self, x_mm: int, y_mm: int) -> Tuple[float, float]:
        """Map coordinate radar -> canvas pixel.
        Radar goc toa do: radar o giua canh duoi. Tang Y la huong vao phia truoc.
        """
        x_ratio = (x_mm - self.ROOM_X_MIN) / (self.ROOM_X_MAX - self.ROOM_X_MIN)
        # Dao Y: canvas (0,0) o goc tren, room Y=0 o goc duoi (radar)
        y_ratio = 1.0 - (y_mm - self.ROOM_Y_MIN) / (self.ROOM_Y_MAX - self.ROOM_Y_MIN)
        # De lai margin
        margin_x = 60
        margin_y = 60
        px = margin_x + x_ratio * (self.CANVAS_W - 2 * margin_x)
        py = margin_y + y_ratio * (self.CANVAS_H - 2 * margin_y)
        return px, py

    def _draw_background(self):
        c = self.canvas
        # Grid lines every 1m
        for x_mm in range(self.ROOM_X_MIN, self.ROOM_X_MAX + 1, 1000):
            x1, y1 = self._mm_to_pixel(x_mm, self.ROOM_Y_MIN)
            x2, y2 = self._mm_to_pixel(x_mm, self.ROOM_Y_MAX)
            c.create_line(x1, y1, x2, y2, fill=self.GRID_COLOR)
            if x_mm != 0:
                c.create_text(x1, self.CANVAS_H - 30, text=f"{x_mm//1000}m", fill="#6c7086", font=("Consolas", 8))
        for y_mm in range(self.ROOM_Y_MIN, self.ROOM_Y_MAX + 1, 1000):
            x1, y1 = self._mm_to_pixel(self.ROOM_X_MIN, y_mm)
            x2, y2 = self._mm_to_pixel(self.ROOM_X_MAX, y_mm)
            c.create_line(x1, y1, x2, y2, fill=self.GRID_COLOR)
            c.create_text(25, y1, text=f"{y_mm//1000}m", fill="#6c7086", font=("Consolas", 8))

        # Radar position (center bottom edge)
        rx, ry = self._mm_to_pixel(0, 0)
        c.create_text(rx, ry + 25, text="RADAR", fill="#f38ba8", font=("Consolas", 10, "bold"))
        c.create_oval(rx - 8, ry - 8, rx + 8, ry + 8, fill="#f38ba8", outline="#f38ba8")

        # Y axis line
        ax, ay = self._mm_to_pixel(0, self.ROOM_Y_MAX)
        c.create_line(rx, ry, ax, ay, fill=self.AXIS_COLOR, dash=(4, 2))

    def _draw_zones(self, zones: List[Zone]):
        c = self.canvas
        for i, z in enumerate(zones):
            x1, y1 = self._mm_to_pixel(z.x1, z.y1)
            x2, y2 = self._mm_to_pixel(z.x2, z.y2)
            # Tkinter can x1<=x2, y1<=y2
            px1, px2 = sorted([x1, x2])
            py1, py2 = sorted([y1, y2])
            color = self.ZONE_OCCUPIED_COLORS[i] if z.occupied else self.ZONE_COLORS[i]
            stipple = "gray25" if not z.occupied else ""
            c.create_rectangle(
                px1, py1, px2, py2,
                outline=color, width=2, fill=color, stipple=stipple
            )
            # Zone label
            lx = (px1 + px2) / 2
            ly = (py1 + py2) / 2
            status = "CO" if z.occupied else "TRONG"
            c.create_text(
                lx, ly,
                text=f"Zone{z.id}\n{status}",
                fill="#1e1e2e" if z.occupied else "#45475a",
                font=("Consolas", 10, "bold")
            )

    def _draw_targets(self, targets: List[Target]):
        c = self.canvas
        now = time.time()
        for i, t in enumerate(targets):
            if not t.valid:
                # Show faded marker if recently seen
                if now - t.last_seen < 3.0:
                    px, py = self._mm_to_pixel(t.x_mm, t.y_mm)
                    c.create_oval(px - 5, py - 5, px + 5, py + 5, outline="#585b70", dash=(2, 2))
                continue

            px, py = self._mm_to_pixel(t.x_mm, t.y_mm)
            color = self.TARGET_COLORS[i]

            # Target dot
            r = 8
            c.create_oval(px - r, py - r, px + r, py + r, fill=color, outline="white", width=2)

            # Velocity vector (speed)
            if t.speed_cm_s != 0:
                scale = 2.0  # pixel per cm/s
                dx = (t.x_mm / max(1, t.distance_mm)) * t.speed_cm_s * scale if t.distance_mm else 0
                dy = (t.y_mm / max(1, t.distance_mm)) * t.speed_cm_s * scale if t.distance_mm else 0
                c.create_line(px, py, px + dx, py - dy, fill=color, width=2, arrow=self.tk.LAST)

            # Label
            label = f"T{t.slot}\n{abs(t.x_mm)}, {t.y_mm}mm\n{abs(t.speed_cm_s)}cm/s"
            c.create_text(
                px + 15, py - 25, anchor="w",
                text=label, fill="white",
                font=("Consolas", 9, "bold"),
                tags=f"tlabel{t.slot}"
            )

    def draw(self):
        self.canvas.delete("all")
        self._draw_background()
        self._draw_zones(self.state.zones)
        self._draw_targets(self.state.targets)

        s = self.state
        self.header_var.set(
            f"Node: {s.node_id}  |  Sequence: {s.sequence}  |  Targets: {s.target_count}/3"
        )
        self.status_var.set(
            f"Frames received: {s.frames_received}   |   "
            f"Last update: {time.time() - s.last_update:.1f}s ago"
        )

    def _on_close(self):
        self._running = False
        self.root.destroy()

    def should_run(self) -> bool:
        return self._running

    def update(self):
        """Call periodically from main loop."""
        if self._running:
            self.draw()
            self.root.update_idletasks()
            self.root.update()


# ============================================================
# Console Headless Mode
# ============================================================
class ConsoleVisualizer:
    def __init__(self, state: RadarState):
        self.state = state
        self._last_print = 0.0

    def update(self):
        now = time.time()
        if now - self._last_print < 0.5:
            return
        self._last_print = now

        s = self.state
        print("\033[2J\033[H", end="")  # Clear screen
        print("=" * 60)
        print(f"  ESP32 LD2453 — Node: {s.node_id}   Seq: {s.sequence}")
        print("=" * 60)
        print(f"  Frames received: {s.frames_received}")
        print(f"  Target count    : {s.target_count}/3")
        if s.last_update:
            print(f"  Last update     : {now - s.last_update:.1f}s ago")
        print("-" * 60)
        for t in s.targets:
            if t.valid:
                status = "DETECTED"
            elif now - t.last_seen < 3.0:
                status = f"RECENT ({now-t.last_seen:.0f}s)"
            else:
                status = "FREE"
            print(f"  Slot {t.slot}: {status}")
            if t.valid:
                print(f"    X={t.x_mm:5d}mm  Y={t.y_mm:5d}mm  "
                      f"Speed={t.speed_cm_s:+3d}cm/s  "
                      f"Dist={int(t.distance_mm):5d}mm  "
                      f"Angle={t.angle_deg:+.1f}°")
        print("-" * 60)
        for z in s.zones:
            occ = "X CO" if z.occupied else "  TRONG"
            print(f"  Zone {z.id}: {occ}  [{z.x1:+5d},{z.y1:+5d}] -> [{z.x2:+5d},{z.y2:+5d}]")
        print("=" * 60)

    def should_run(self) -> bool:
        return True

    def draw(self):
        pass


# ============================================================
# Fetch zones from ESP32 REST API
# ============================================================
async def fetch_zones(state: RadarState, esp32_host: str):
    try:
        import aiohttp
    except ImportError:
        print("Canh bao: aiohttp chua cai. Zone se dung default. (pip install aiohttp)")
        return

    url = f"http://{esp32_host}/zones"
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=3)) as session:
            async with session.get(url) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    for z_data in data:
                        idx = z_data.get("id", 0)
                        if 1 <= idx <= 3:
                            z = state.zones[idx - 1]
                            z.id = idx
                            z.x1 = int(z_data.get("x1", z.x1))
                            z.y1 = int(z_data.get("y1", z.y1))
                            z.x2 = int(z_data.get("x2", z.x2))
                            z.y2 = int(z_data.get("y2", z.y2))
                            z.occupied = bool(z_data.get("occupied", False))
                    print(f"Da lay {len(data)} zone tu ESP32")
                else:
                    print(f"GET zones tra ve status {resp.status}")
    except Exception as e:
        print(f"Khong the lay zones: {e}")


# ============================================================
# Update zone occupancy from target positions
# (redundant — ESP32 already reports zones, but kept for validation)
# ============================================================
def update_zone_occupancy_from_targets(state: RadarState):
    for z in state.zones:
        z.occupied = False
        for t in state.targets:
            if not t.valid:
                continue
            if z.x1 <= t.x_mm <= z.x2 and z.y1 <= t.y_mm <= z.y2:
                z.occupied = True
                break


# ============================================================
# Main async loop
# ============================================================
async def ws_client_loop(esp32_host: str, visualizer, state: RadarState):
    uri = f"ws://{esp32_host}/ws"
    reconnect_delay = 2.0

    while visualizer.should_run():
        print(f"Dang ket noi den {uri}...")
        try:
            async with websockets.connect(uri, ping_interval=30, ping_timeout=10) as ws:
                print(f"Da ket noi WebSocket den ESP32 ({esp32_host})")
                # Fetch zones one time after connect
                await fetch_zones(state, esp32_host)

                async for message in ws:
                    if not visualizer.should_run():
                        break
                    try:
                        frame = json.loads(message)
                        state.update_from_frame(frame)
                        update_zone_occupancy_from_targets(state)
                        visualizer.update()
                    except json.JSONDecodeError:
                        print(f"Loi parse JSON: {message[:80]}...")
                        continue
                # Graceful disconnect
                print("WebSocket dong. Dang ket noi lai...")
        except websockets.exceptions.ConnectionClosed as e:
            print(f"Mat ket noi WebSocket ({e}), dang thu lai sau {reconnect_delay}s...")
            await asyncio.sleep(reconnect_delay)
        except OSError as e:
            print(f"Khong the ket noi: {e}. Thu lai sau {reconnect_delay}s...")
            await asyncio.sleep(reconnect_delay)
        except KeyboardInterrupt:
            print("Dung lai...")
            break
        except Exception as e:
            print(f"Loi WebSocket: {type(e).__name__}: {e}")
            await asyncio.sleep(reconnect_delay)


def main():
    parser = argparse.ArgumentParser(description="ESP32 LD2453 Radar Visualizer")
    parser.add_argument("host", help="ESP32 IP address (e.g. 192.168.1.100)")
    parser.add_argument("--headless", action="store_true", help="Chi in ra console, khong can GUI")
    parser.add_argument("--mode", choices=["tkinter", "pygame"], default="tkinter",
                        help="Loaị GUI (mac dinh: tkinter)")
    args = parser.parse_args()

    state = RadarState()
    # Default zones (giong ESP32)
    state.zones[0] = Zone(1, x1=-4000, y1=1, x2=-1, y2=4000)
    state.zones[1] = Zone(2, x1=1, y1=1, x2=4000, y2=4000)
    state.zones[2] = Zone(3, x1=-4001, y1=4001, x2=4001, y2=6000)

    if args.headless:
        vis = ConsoleVisualizer(state)
    elif args.mode == "pygame":
        print("Pygame mode chua implement, chuyen sang tkinter...")
        vis = TkinterVisualizer(state)
    else:
        vis = TkinterVisualizer(state)

    print(f"ESP32-S3 LD2453 Radar Visualizer")
    print(f"ESP32 host: {args.host}")
    print(f"Mode: {'headless console' if args.headless else 'tkinter GUI'}")

    try:
        asyncio.run(ws_client_loop(args.host, vis, state))
    except KeyboardInterrupt:
        print("\nThoat.")


if __name__ == "__main__":
    main()