#!/usr/bin/env python3
"""
ESP32-S3 LD2453 Radar — Web Visualizer for Orange Pi
======================================================
- Kết nối WebSocket tới ESP32 nhận dữ liệu radar realtime
- Mở Web Server (FastAPI) cho phép trình duyệt xem đồ họa đẹp
- Broadcast dữ liệu từ ESP32 tới tất cả các client trình duyệt đang kết nối
- Hỗ trợ REST API lấy zones, cập nhật zones (proxy tới ESP32)

Sử dụng:
  python visualizer.py <ESP32_IP>                    # Mặc định port 8000
  python visualizer.py 192.168.1.100 --port 8080     # Tùy chỉnh port
  python visualizer.py 192.168.1.100 --host 0.0.0.0  # Mở cho truy cập ngoài

Mở trình duyệt vào: http://<OrangePi_IP>:8000
"""

import argparse
import asyncio
import json
import logging
import math
import sys
import time
from contextlib import asynccontextmanager
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import List, Set

# ------------------------------------------------------------
# Third-party imports (install via requirements.txt)
# ------------------------------------------------------------
try:
    import websockets
    from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
    from fastapi.responses import FileResponse
    from fastapi.staticfiles import StaticFiles
    import aiohttp
except ImportError as e:
    print(f"Lỗi: Thiếu thư viện. Cần cài đặt trước: pip install -r requirements.txt")
    print(f"Chi tiết: {e}")
    sys.exit(1)


# ============================================================
# Data structures — keep same semantics as original
# ============================================================
@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
    esp32_connected: bool = False

    def to_json(self) -> dict:
        """Serialize to something the browser can render."""
        now = time.time()
        return {
            "node_id": self.node_id,
            "sequence": self.sequence,
            "timestamp_ms": self.timestamp_ms,
            "target_count": self.target_count,
            "targets": [
                {
                    "slot": t.slot,
                    "valid": t.valid,
                    "x_mm": t.x_mm,
                    "y_mm": t.y_mm,
                    "speed_cm_s": t.speed_cm_s,
                    "distance_mm": round(t.distance_mm, 1),
                    "angle_deg": round(t.angle_deg, 2),
                    # seconds since this target was last valid; None if never seen
                    "last_seen": (
                        round(now - t.last_seen, 2)
                        if t.valid or (t.last_seen and t.last_seen > 0)
                        else None
                    ),
                }
                for t in self.targets
            ],
            "zones": [
                {
                    "id": z.id,
                    "x1": z.x1,
                    "y1": z.y1,
                    "x2": z.x2,
                    "y2": z.y2,
                    "occupied": z.occupied,
                }
                for z in self.zones
            ],
            "frames_received": self.frames_received,
            "last_update": (
                round(now - self.last_update, 2)
                if self.last_update and self.last_update > 0
                else None
            ),
            "esp32_connected": self.esp32_connected,
            "server_time": round(now, 2),
        }

    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()

        # Belt-and-suspenders physical-limits filter.  The ESP32 side
        # should already have gated these, but in case old firmware is
        # still running or an ACK frame leaks through, we re-apply the
        # same sanity rules on the Orange Pi side:
        #   Y in front of radar (y >= 1 mm), <= 6000 mm rated range
        #   |X| <= 6000 mm (±60° cone + safety margin)
        PHYS_Y_MIN = 1
        PHYS_Y_MAX = 6000
        PHYS_X_ABS = 6000
        valid_targets_this_frame = 0

        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 = bool(t_data.get("valid", False))
                t.x_mm = int(t_data.get("x", 0))
                t.y_mm = int(t_data.get("y", 0))
                t.speed_cm_s = int(t_data.get("speed", 0))

                # If ESP32 parser already gives distance/angle, use them
                t.distance_mm = float(t_data.get("distance", 0.0))
                t.angle_deg = float(t_data.get("angle", 0.0))

                # Sanity overrides: strip validity if coordinates lie
                # outside the radar's measurable physical envelope.
                y_ok = PHYS_Y_MIN <= t.y_mm <= PHYS_Y_MAX
                x_ok = -PHYS_X_ABS <= t.x_mm <= PHYS_X_ABS
                if not (y_ok and x_ok):
                    t.valid = False

                # Derive distance / angle if missing, for convenience
                if t.valid and (t.distance_mm == 0 or t.angle_deg == 0):
                    t.distance_mm = math.hypot(t.x_mm, t.y_mm)
                    if t.distance_mm > 0:
                        t.angle_deg = math.degrees(math.atan2(t.x_mm, t.y_mm))

                if t.valid:
                    t.last_seen = now
                    valid_targets_this_frame += 1

        # Keep the browser-side counters honest even if the ESP32 told
        # us a target_count that included the now-filtered ACK ghost.
        self.target_count = valid_targets_this_frame
        self.frames_received += 1
        self.last_update = now


# ============================================================
# Global app state
# ============================================================
STATE = RadarState()
# Default zones — matches the original ESP32 layout
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)

CONNECTED_BROWSERS: Set[WebSocket] = set()
BROWSER_LOCK = asyncio.Lock()
ESP32_HOST: str = ""  # populated from CLI

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
log = logging.getLogger("radar-web")


# ============================================================
# Helper: update zone occupancy (redundant but useful 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 (min(z.x1, z.x2) <= t.x_mm <= max(z.x1, z.x2)
                    and min(z.y1, z.y2) <= t.y_mm <= max(z.y1, z.y2)):
                z.occupied = True
                break


# ============================================================
# Heartbeat: broadcast last known state to browsers every 1s
# even if ESP32 hasn't sent a new frame yet (parser stub scenario).
# ============================================================
async def heartbeat_broadcast():
    while True:
        try:
            await broadcast_to_browsers()
        except Exception:
            pass
        await asyncio.sleep(1.0)


# ============================================================
# ESP32 WebSocket client — runs as a background task
# ============================================================
async def esp32_ws_forever(esp32_host: str):
    """Loop forever: connect to ESP32 /ws, read frames, update STATE, broadcast."""
    uri = f"ws://{esp32_host}/ws"
    reconnect_delay = 2.0

    while True:
        try:
            STATE.esp32_connected = False
            log.info(f"Đang kết nối ESP32 tại {uri} ...")
            # Push initial "not connected" state to any browser that joined early
            await broadcast_to_browsers()

            async with websockets.connect(
                uri, ping_interval=30, ping_timeout=10, close_timeout=5,
                open_timeout=8, max_size=1_000_000,
            ) as ws:
                STATE.esp32_connected = True
                log.info(f"Đã kết nối WebSocket tới ESP32 ({esp32_host})")

                # Fetch zones from ESP32 after successful connect
                await fetch_zones_from_esp32(esp32_host)
                # Broadcast immediately so browsers see zones right after connect
                await broadcast_to_browsers()

                async for message in ws:
                    try:
                        frame = json.loads(message)
                    except json.JSONDecodeError:
                        log.warning(f"Lỗi parse JSON từ ESP32: {message[:80]}...")
                        continue

                    STATE.update_from_frame(frame)
                    _update_zone_occupancy_from_targets(STATE)
                    # Always push the latest state
                    await broadcast_to_browsers()

            log.info("ESP32 WebSocket đóng, đang kết nối lại...")
            STATE.esp32_connected = False
            await broadcast_to_browsers()

        except websockets.exceptions.ConnectionClosed as e:
            STATE.esp32_connected = False
            log.warning(f"Mất kết nối ESP32 ({e}), thử lại sau {reconnect_delay}s")
            await broadcast_to_browsers()
            await asyncio.sleep(reconnect_delay)
        except OSError as e:
            STATE.esp32_connected = False
            log.warning(f"Không thể kết nối ESP32: {e}. Thử lại sau {reconnect_delay}s")
            await broadcast_to_browsers()
            await asyncio.sleep(reconnect_delay)
        except asyncio.CancelledError:
            log.info("Task ESP32 WS bị hủy (shutdown).")
            break
        except Exception as e:
            STATE.esp32_connected = False
            log.exception(f"Lỗi ESP32 WS loop: {type(e).__name__}: {e}")
            await asyncio.sleep(reconnect_delay)


async def fetch_zones_from_esp32(esp32_host: str):
    """Fetch zone layout from ESP32 /zones REST endpoint."""
    url = f"http://{esp32_host}/zones"
    try:
        timeout = aiohttp.ClientTimeout(total=3)
        async with aiohttp.ClientSession(timeout=timeout) 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))
                    log.info(f"Đã lấy {len(data)} zones từ ESP32")
                else:
                    log.warning(f"GET /zones trả về HTTP {resp.status}")
    except Exception as e:
        log.warning(f"Không lấy được zones từ ESP32 (sẽ dùng mặc định): {e}")


async def post_zones_to_esp32(esp32_host: str, zones_payload):
    """Forward a zone update to ESP32 /updateZones."""
    url = f"http://{esp32_host}/updateZones"
    try:
        timeout = aiohttp.ClientTimeout(total=5)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(url, json=zones_payload) as resp:
                body = await resp.text()
                if resp.status in (200, 204):
                    # Apply locally immediately
                    for z_data in zones_payload:
                        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))
                    await broadcast_to_browsers()
                    return {"status": "ok", "http_status": resp.status, "body": body}
                else:
                    raise HTTPException(status_code=resp.status, detail=body)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Lỗi khi gọi ESP32: {e}")


# ============================================================
# Broadcast to all connected browsers
# ============================================================
async def broadcast_to_browsers():
    payload = json.dumps(STATE.to_json())
    if not CONNECTED_BROWSERS:
        return
    async with BROWSER_LOCK:
        dead: List[WebSocket] = []
        for ws in list(CONNECTED_BROWSERS):
            try:
                await ws.send_text(payload)
            except Exception:
                dead.append(ws)
        for ws in dead:
            CONNECTED_BROWSERS.discard(ws)


# ============================================================
# FastAPI app setup
# ============================================================
BASE_DIR = Path(__file__).resolve().parent
TEMPLATES_DIR = BASE_DIR / "templates"
STATIC_DIR = BASE_DIR / "static"
STATIC_DIR.mkdir(exist_ok=True)
TEMPLATES_DIR.mkdir(exist_ok=True)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Start background tasks when app starts, cancel them on shutdown."""
    tasks = [
        asyncio.create_task(esp32_ws_forever(ESP32_HOST)),
        asyncio.create_task(heartbeat_broadcast()),
    ]
    yield
    for t in tasks:
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)


app = FastAPI(
    title="ESP32 LD2453 Radar Web Visualizer",
    lifespan=lifespan,
)

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


# --- HTTP routes ------------------------------------------------
@app.get("/")
async def get_index():
    index_file = TEMPLATES_DIR / "index.html"
    if not index_file.exists():
        return {
            "error": "index.html không tồn tại. Hãy đảm bảo thư mục templates/ có index.html."
        }
    return FileResponse(str(index_file))


@app.get("/api/state")
async def get_state():
    """REST endpoint lấy trạng thái hiện tại (polling dự phòng)."""
    return STATE.to_json()


@app.get("/api/zones")
async def get_zones():
    return [
        {"id": z.id, "x1": z.x1, "y1": z.y1, "x2": z.x2, "y2": z.y2, "occupied": z.occupied}
        for z in STATE.zones
    ]


@app.post("/api/zones")
async def update_zones(payload: list):
    """Nhận cập nhật zones từ trình duyệt và forward tới ESP32."""
    if not ESP32_HOST:
        raise HTTPException(status_code=500, detail="ESP32 host chưa được cấu hình")
    return await post_zones_to_esp32(ESP32_HOST, payload)


# --- WebSocket route for browsers -----------------------------
@app.websocket("/ws")
async def browser_ws_endpoint(ws: WebSocket):
    await ws.accept()
    # Gửi trạng thái ngay lập tức để client có data lần đầu
    try:
        await ws.send_text(json.dumps(STATE.to_json()))
    except Exception:
        return

    async with BROWSER_LOCK:
        CONNECTED_BROWSERS.add(ws)
    log.info(f"Browser kết nối: {ws.client} (tổng: {len(CONNECTED_BROWSERS)})")

    try:
        # Browser chủ yếu nhận, nhưng chúng ta vẫn đọc để phát hiện ngắt kết nối
        while True:
            data = await ws.receive_text()
            # Có thể mở rộng: nhận lệnh từ trình duyệt (ví dụ: yêu cầu cập nhật zones)
            try:
                msg = json.loads(data)
                if msg.get("type") == "updateZones":
                    zones = msg.get("zones", [])
                    await post_zones_to_esp32(ESP32_HOST, zones)
            except Exception:
                pass
    except WebSocketDisconnect:
        pass
    finally:
        async with BROWSER_LOCK:
            CONNECTED_BROWSERS.discard(ws)
        log.info(f"Browser ngắt kết nối (còn lại: {len(CONNECTED_BROWSERS)})")


# ============================================================
# Entry point
# ============================================================
def main():
    global ESP32_HOST

    parser = argparse.ArgumentParser(
        description="ESP32 LD2453 Radar — Web Visualizer (Xem trên trình duyệt)"
    )
    parser.add_argument(
        "esp32",
        help="Địa chỉ IP ESP32 (ví dụ: 192.168.1.100)",
    )
    parser.add_argument(
        "--host",
        default="0.0.0.0",
        help="Host lắng nghe của Web Server (mặc định: 0.0.0.0 — cho phép truy cập từ ngoài)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port của Web Server (mặc định: 8000)",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Hiển thị log chi tiết (DEBUG level)",
    )
    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    ESP32_HOST = args.esp32.strip()

    # Try importing uvicorn — keep it here so help text still shows without install
    try:
        import uvicorn
    except ImportError:
        print("Lỗi: Cần cài uvicorn: pip install uvicorn")
        sys.exit(1)

    print()
    print("=" * 62)
    print("  ESP32-S3 LD2453 Radar — Web Visualizer")
    print("=" * 62)
    print(f"  ESP32   : ws://{ESP32_HOST}/ws")
    print(f"  Web UI  : http://{args.host}:{args.port}/")
    print(f"  WS UI   : ws://<OrangePi_IP>:{args.port}/ws")
    print(f"  Browsers: 0 client đang kết nối")
    print("=" * 62)
    print("  Nhấn Ctrl+C để dừng")
    print("=" * 62)
    print()

    try:
        uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
    except KeyboardInterrupt:
        print("\nThoát.")


if __name__ == "__main__":
    main()