# WebSocket client — main event loop, service request/response, DeviceBridge
import asyncio
import base64
import json
import threading
import time
import uuid

from .config import (
    log, _ros2_cache, _ros2_lock, _gw_context,
    _service_futures,
    DEVICE_ID, GATEWAY_URL, API_KEY, MODE, INTERVAL,
)
import ai_teammate_ros2_bridge.config as _cfg
from .commands import handle_command
from .ros2_node import _ros2_spin_thread
from .sensors import get_ros2_data, get_sim_data
from .keyframe import _should_capture_keyframe, _capture_keyframe, _maybe_upload_keyframe_batch


# ── Service request/response ──

async def _send_service_request(data: dict, timeout: float = 15.0) -> dict:
    """Send service_request WS message and wait for service_response."""
    request_id = str(uuid.uuid4())[:8]
    data["request_id"] = request_id

    future = asyncio.get_event_loop().create_future()
    _service_futures[request_id] = future

    try:
        await _cfg._ws_ref.send(json.dumps({"type": "service_request", "data": data}))
        result = await asyncio.wait_for(future, timeout=timeout)
        return result
    except asyncio.TimeoutError:
        return {"status": "error", "message": "service request timeout"}
    except Exception as e:
        return {"status": "error", "message": f"service request failed: {e}"}
    finally:
        _service_futures.pop(request_id, None)


def _handle_service_response(data: dict):
    """Route service_response to pending future."""
    request_id = data.get("request_id", "")
    future = _service_futures.get(request_id)
    if future and not future.done():
        future.set_result(data)


# ── DeviceBridge class ──

class DeviceBridge:
    """Main device bridge — connects ROS2 robot to AI Teammate Gateway."""

    def __init__(self, device_id: str, gateway_url: str, api_key: str, mode: str = "rclpy"):
        self.device_id = device_id
        self.gateway_url = gateway_url
        self.api_key = api_key
        self.mode = mode

    @classmethod
    def from_env(cls):
        return cls(
            device_id=DEVICE_ID,
            gateway_url=GATEWAY_URL,
            api_key=API_KEY,
            mode=MODE,
        )

    def run(self):
        log.info(f"Bridge v2 starting: device={self.device_id} mode={self.mode}")
        asyncio.run(self._run())

    async def _run(self):
        import websockets

        url = f"{self.gateway_url}/ws/{self.device_id}"
        tick = 0

        # Start rclpy spin thread
        if self.mode in ("rclpy", "rosbridge"):
            t = threading.Thread(target=_ros2_spin_thread, daemon=True)
            t.start()
            log.info("rclpy spin thread started, waiting for node...")
            for _ in range(50):
                if _cfg._cmd_vel_pub is not None:
                    break
                await asyncio.sleep(0.1)
            if _cfg._cmd_vel_pub:
                log.info("rclpy node ready")
            else:
                log.warning("rclpy node not ready after 5s, continuing anyway")

        while True:
            try:
                async with websockets.connect(url, ping_interval=20, ping_timeout=10, max_size=2**20) as ws:
                    log.info(f"Connected to {self.gateway_url} as {self.device_id} (mode={self.mode})")

                    # Auth
                    await ws.send(json.dumps({"type": "auth", "data": {"api_key": self.api_key}}))
                    log.info("Auth message sent")

                    try:
                        welcome = await asyncio.wait_for(ws.recv(), timeout=5)
                        log.info(f"Welcome received: {str(welcome)[:200]}")
                        try:
                            wdata = json.loads(welcome) if isinstance(welcome, str) else welcome
                            ctx = wdata.get("data", {}).get("context", {})
                            if ctx:
                                _gw_context.update(ctx)
                                log.info(f"Gateway context: {_gw_context}")
                        except (json.JSONDecodeError, AttributeError):
                            pass
                    except asyncio.TimeoutError:
                        log.warning("Welcome message timeout")

                    # Apply calibration from context if available
                    cal = _gw_context.get("calibration")
                    if cal and isinstance(cal, dict):
                        if cal.get("rotation_scale") and 0.5 <= cal["rotation_scale"] <= 2.0:
                            _cfg.ROTATION_SCALE = cal["rotation_scale"]
                        if cal.get("distance_scale") and 0.5 <= cal["distance_scale"] <= 2.0:
                            _cfg.DISTANCE_SCALE = cal["distance_scale"]
                        log.info(f"Calibration loaded: rot={_cfg.ROTATION_SCALE}, dist={_cfg.DISTANCE_SCALE}")

                    # Set WS ref
                    _cfg._ws_ref = ws

                    _LONG_COMMANDS = {
                        "MOVE", "TURN", "SAFE_MOVE", "SAFE_NAVIGATE_PATH", "SAFE_CIRCLE", "CIRCLE",
                        "NAVIGATE_PATH", "PATROL", "SLAM_SAVE", "NAV2_GOAL", "GO_TO_LOCATION", "DOCK", "UNDOCK",
                        "KEYFRAME_START", "KEYFRAME_STOP", "MAP_UPLOAD", "MAP_DOWNLOAD",
                        "CALIBRATE_ROTATION", "CALIBRATE_DISTANCE", "CALIBRATE_FULL",
                    }

                    async def _run_and_respond(message):
                        try:
                            result = await handle_command(message)
                            if result:
                                resp = json.dumps({
                                    "type": "command_response",
                                    "data": {
                                        "command_id": message.get("data", {}).get("command_id", ""),
                                        **result,
                                    },
                                })
                                await ws.send(resp)
                                log.info(f"Command response sent: {result.get('status')}")
                        except Exception as e:
                            log.error(f"Command execution error: {e}")

                    async def on_recv(raw):
                        log.info(f"WS RECV: {str(raw)[:120]}")
                        try:
                            message = json.loads(raw)
                            msg_type = message.get("type", "")

                            if msg_type in ("pong",):
                                return

                            if msg_type == "service_response":
                                _handle_service_response(message.get("data", {}))
                                return

                            cmd_type = message.get("data", {}).get("command_type", "")

                            if cmd_type in _LONG_COMMANDS:
                                asyncio.create_task(_run_and_respond(message))
                            else:
                                await _run_and_respond(message)
                        except json.JSONDecodeError:
                            pass

                    async def recv_loop():
                        try:
                            async for raw in ws:
                                await on_recv(raw)
                        except Exception as e:
                            log.warning(f"recv_loop error: {e}")

                    recv_task = asyncio.create_task(recv_loop())

                    while True:
                        if self.mode in ("rclpy", "rosbridge"):
                            data = get_ros2_data()
                            if not data:
                                data = get_sim_data(tick)
                        else:
                            data = get_sim_data(tick)

                        msg = json.dumps({
                            "type": "status_update",
                            "data": {"device_id": self.device_id, "timestamp": int(time.time() * 1000), **data},
                        })
                        await ws.send(msg)

                        # Camera stream
                        with _ros2_lock:
                            streaming = _ros2_cache.get("camera_stream", False)
                            stream_fps = _ros2_cache.get("camera_stream_fps", 2)
                            frame = _ros2_cache.get("camera_frame") if streaming else None
                        if frame and tick % max(1, int(1.0 / INTERVAL / stream_fps)) == 0:
                            await ws.send(json.dumps({
                                "type": "frame_data",
                                "data": {
                                    "device_id": self.device_id,
                                    "image_base64": base64.b64encode(frame).decode('ascii'),
                                    "format": "jpeg",
                                    "timestamp": int(time.time() * 1000),
                                },
                            }))

                        # Keyframe capture
                        if _cfg._kf_capture_active and _should_capture_keyframe():
                            _capture_keyframe()
                            await _maybe_upload_keyframe_batch()

                        if tick % 10 == 0:
                            log.info(f"Sent {tick} status_updates")
                        tick += 1
                        await asyncio.sleep(INTERVAL)

                        if recv_task.done():
                            break

            except Exception as e:
                log.warning(f"Disconnected: {e}, reconnecting in 5s...")
                await asyncio.sleep(5)
