"""DeviceAgent — the main orchestrator tying ACP, MQTT, WebRTC, and AI together.

A DeviceAgent represents a single device (1 device = 1 process = 1 BeeOS Instance).
It manages three independent channels:
  1. ACP (via Bridge) — control plane: AI chat + device management
  2. MQTT (via EMQX)  — signaling + telemetry
  3. WebRTC (P2P/TURN) — on-demand video streaming + DataChannel control
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Any

from beeos_bridge_client import BridgeClient, BridgeClientOptions, ReconnectConfig

from .acp.device_handler import DeviceHandler
from .acp.router import ACPRouter
from .acp.session_handler import SessionHandler
from .ai.mobile_agent_bridge import AgentStep, MobileAgentBridge
from .ai.task_engine import TaskEngine
from .mqtt.client import MQTTManager
from .mqtt.signaling import SignalingHandler
from .mqtt.stream_lifecycle import StreamLifecycle
from .mqtt.telemetry import TelemetryReporter
from .runtime.base import DeviceRuntime
from .webrtc.peer import WebRTCPeer
from .webrtc.scrcpy_adapter import ScrcpyAdapter

logger = logging.getLogger(__name__)


@dataclass
class DeviceAgentConfig:
    bridge_url: str = ""
    bridge_service: str = "acp"
    standalone: bool = False
    key_file: str | None = None
    private_key: bytes | None = None
    public_key: bytes | None = None

    mqtt_host: str = ""
    mqtt_port: int = 8883
    mqtt_username: str | None = None
    mqtt_password: str | None = None
    mqtt_use_tls: bool = True

    device_id: str = ""

    vlm_api_key: str = ""
    vlm_base_url: str = ""
    vlm_model: str = ""

    ice_servers: list[dict[str, Any]] = field(default_factory=list)
    cooldown_seconds: float = 30.0
    telemetry_interval: float = 30.0

    http_enabled: bool = False
    http_host: str = "0.0.0.0"
    http_port: int = 9090


class DeviceAgent:
    """Main DeviceAgent orchestrating three channels + AI automation."""

    def __init__(
        self,
        config: DeviceAgentConfig,
        runtime: DeviceRuntime,
    ) -> None:
        self._config = config
        self._runtime = runtime
        self._running = False

        self._bridge: BridgeClient | None = None
        self._mqtt: MQTTManager | None = None
        self._webrtc_peer: WebRTCPeer | None = None
        self._scrcpy: ScrcpyAdapter | None = None

        self._acp_router: ACPRouter | None = None
        self._signaling: SignalingHandler | None = None
        self._stream_lifecycle: StreamLifecycle | None = None
        self._telemetry: TelemetryReporter | None = None
        self._ai_bridge: MobileAgentBridge | None = None
        self._task_engine: TaskEngine | None = None
        self._http_api = None
        self._touch_state: dict[int, tuple[int, int]] = {}  # pointerId → (startX, startY)

    @property
    def running(self) -> bool:
        return self._running

    async def start(self) -> None:
        """Initialize all channels and start the agent.

        In standalone mode only the device runtime, task engine, and HTTP API
        are started — Bridge (ACP) and MQTT are skipped entirely.
        """
        standalone = self._config.standalone
        logger.info(
            "Starting DeviceAgent%s...", " (standalone)" if standalone else ""
        )

        await self._runtime.connect()
        info = await self._runtime.get_info()
        logger.info("Device: %s (%s, %dx%d)", info.name, info.os, info.width, info.height)

        device_id = self._config.device_id or info.metadata.get("serial", "unknown")

        if not standalone:
            # Channel 1: ACP (Bridge)
            await self._start_acp(device_id)

            # Channel 2: MQTT (EMQX) — only if configured
            if self._config.mqtt_host:
                await self._start_mqtt(device_id)

        self._task_engine = TaskEngine(
            runtime=self._runtime,
            on_step=self._on_task_step,
            vlm_api_key=self._config.vlm_api_key,
            vlm_base_url=self._config.vlm_base_url,
            vlm_model=self._config.vlm_model,
        )
        await self._task_engine.start()

        # HTTP API — always on in standalone mode, otherwise opt-in
        if self._config.http_enabled or standalone:
            await self._start_http_api()

        self._running = True
        logger.info("DeviceAgent started (device_id=%s)", device_id)

    async def _start_acp(self, device_id: str) -> None:
        """Initialize the ACP channel via Bridge."""
        opts = BridgeClientOptions(
            bridge_url=self._config.bridge_url,
            service=self._config.bridge_service,
            key_file=self._config.key_file,
            private_key=self._config.private_key,
            public_key=self._config.public_key,
            reconnect=ReconnectConfig(enabled=True),
        )
        self._bridge = BridgeClient(opts)

        device_handler = DeviceHandler(self._runtime)
        session_handler = SessionHandler(on_prompt=self._on_ai_prompt)
        self._acp_router = ACPRouter(device_handler, session_handler)

        self._bridge.on("message", self._on_acp_message)
        self._bridge.on("connected", self._on_acp_connected)
        self._bridge.on("disconnected", lambda r: logger.warning("ACP channel disconnected: %s", r))

        await self._bridge.connect()

    async def _start_mqtt(self, device_id: str) -> None:
        """Initialize the MQTT channel for signaling and telemetry."""
        self._mqtt = MQTTManager(
            device_id=device_id,
            mqtt_host=self._config.mqtt_host,
            mqtt_port=self._config.mqtt_port,
            mqtt_username=self._config.mqtt_username,
            mqtt_password=self._config.mqtt_password,
            use_tls=self._config.mqtt_use_tls,
        )

        self._stream_lifecycle = StreamLifecycle(
            cooldown_seconds=self._config.cooldown_seconds,
            on_start_stream=self._on_start_stream,
            on_stop_stream=self._on_stop_stream,
        )

        self._webrtc_peer = WebRTCPeer(
            ice_servers=self._config.ice_servers,
            on_viewer_connected=self._stream_lifecycle.viewer_connected,
            on_viewer_disconnected=self._stream_lifecycle.viewer_disconnected,
            on_data_channel_message=self._on_datachannel_message,
            on_prepare_track=self._prepare_video_track,
        )

        self._signaling = SignalingHandler(
            mqtt=self._mqtt,
            on_offer=self._on_webrtc_offer,
            on_ice_candidate=self._on_ice_candidate,
        )

        self._telemetry = TelemetryReporter(
            mqtt=self._mqtt,
            interval=self._config.telemetry_interval,
            get_device_metrics=self._get_device_metrics,
        )

        await self._mqtt.connect()
        await self._telemetry.start()

    async def _get_device_metrics(self) -> dict[str, Any]:
        """Collect device battery/CPU metrics via ADB shell for telemetry."""
        metrics: dict[str, Any] = {}
        try:
            battery_raw = await self._runtime.execute_shell("dumpsys battery")
            for line in battery_raw.splitlines():
                line = line.strip()
                if line.startswith("level:"):
                    metrics["battery_level"] = int(line.split(":", 1)[1].strip())
                elif line.startswith("status:"):
                    val = int(line.split(":", 1)[1].strip())
                    metrics["battery_charging"] = val == 2 or val == 5
                elif line.startswith("temperature:"):
                    metrics["battery_temp"] = int(line.split(":", 1)[1].strip()) / 10.0
        except Exception:
            logger.debug("Failed to collect battery metrics", exc_info=True)
        try:
            mem_raw = await self._runtime.execute_shell(
                "cat /proc/meminfo | head -3"
            )
            mem_total = mem_free = 0
            for line in mem_raw.splitlines():
                if line.startswith("MemTotal:"):
                    mem_total = int(line.split()[1])
                elif line.startswith("MemAvailable:"):
                    mem_free = int(line.split()[1])
            if mem_total > 0:
                metrics["memory_total_kb"] = mem_total
                metrics["memory_available_kb"] = mem_free
                metrics["memory_usage_pct"] = round(
                    (1 - mem_free / mem_total) * 100, 1
                )
        except Exception:
            logger.debug("Failed to collect memory metrics", exc_info=True)
        return metrics

    async def _start_http_api(self) -> None:
        from .api import DeviceAgentAPI
        self._http_api = DeviceAgentAPI(
            runtime=self._runtime,
            task_engine=self._task_engine,
            host=self._config.http_host,
            port=self._config.http_port,
        )
        await self._http_api.start()

    async def stop(self) -> None:
        """Gracefully shut down all channels."""
        logger.info("Stopping DeviceAgent...")
        self._running = False

        if self._http_api:
            await self._http_api.stop()

        if self._task_engine:
            await self._task_engine.stop()

        if self._telemetry:
            await self._telemetry.stop()

        if self._stream_lifecycle:
            await self._stream_lifecycle.force_stop()

        if self._webrtc_peer:
            await self._webrtc_peer.close()

        if self._scrcpy and self._scrcpy.running:
            await self._scrcpy.stop()

        if self._mqtt:
            await self._mqtt.disconnect()

        if self._bridge:
            await self._bridge.close()

        await self._runtime.disconnect()
        logger.info("DeviceAgent stopped")

    # -- ACP connection / device registration --

    async def _on_acp_connected(self) -> None:
        logger.info("ACP channel connected")
        asyncio.get_event_loop().call_soon(
            lambda: asyncio.ensure_future(self._report_device_info())
        )

    async def _report_device_info(self) -> None:
        """Send device info notification via ACP after connection."""
        try:
            info = await self._runtime.get_info()
            from dataclasses import asdict
            info_dict = asdict(info)
            info_dict["type"] = info.type.value
            info_dict["capabilities"] = [c.value for c in info.capabilities]

            notification = json.dumps({
                "jsonrpc": "2.0",
                "method": "device/register",
                "params": info_dict,
            })
            if self._bridge:
                await self._bridge.send(notification)
                logger.info("Reported device info via ACP: %s (%s)", info.name, info.os)
        except Exception:
            logger.exception("Failed to report device info")

        # Also publish to MQTT for backend consumption
        if self._mqtt:
            try:
                info = await self._runtime.get_info()
                from dataclasses import asdict
                info_dict = asdict(info)
                info_dict["type"] = info.type.value
                info_dict["capabilities"] = [c.value for c in info.capabilities]
                await self._mqtt.publish("info", info_dict)
                logger.info("Published device info to MQTT")
            except Exception:
                logger.exception("Failed to publish device info to MQTT")

    # -- ACP message handling --

    async def _on_acp_message(self, raw: str | bytes) -> None:
        if isinstance(raw, bytes):
            raw = raw.decode()
        response = await self._acp_router.handle_message(raw)
        if response and self._bridge:
            await self._bridge.send(response)

    # -- AI prompt handling --

    async def _on_ai_prompt(self, data: dict[str, Any]) -> None:
        """Handle session/prompt — submit to TaskEngine."""
        prompt = data.get("prompt", "")
        if not prompt:
            return

        if self._telemetry:
            self._telemetry.agent_status = "running"

        if self._task_engine:
            task = await self._task_engine.submit(prompt)
            if self._bridge:
                await self._bridge.send(json.dumps({
                    "jsonrpc": "2.0",
                    "method": "agent_message_chunk",
                    "params": {
                        "task_id": task.id,
                        "status": "accepted",
                        "instruction": prompt,
                    },
                }))

    async def _on_task_step(self, task_id: str, step: dict[str, Any]) -> None:
        """Stream task step progress via ACP."""
        if self._bridge:
            await self._bridge.send(json.dumps({
                "jsonrpc": "2.0",
                "method": "agent_message_chunk",
                "params": {"task_id": task_id, **step},
            }))

    # -- WebRTC signaling --

    async def _on_webrtc_offer(self, payload: dict[str, Any]) -> dict[str, Any] | None:
        """Handle WebRTC offer from a viewer."""
        if not self._webrtc_peer:
            return None
        return await self._webrtc_peer.handle_offer(payload)

    async def _on_ice_candidate(self, payload: dict[str, Any]) -> None:
        """Handle ICE candidate from a viewer."""
        if self._webrtc_peer:
            await self._webrtc_peer.handle_ice_candidate(payload)

    # -- Stream lifecycle callbacks --

    async def _prepare_video_track(self) -> Any:
        """Start scrcpy and return the video track, called by WebRTCPeer before creating the answer."""
        if self._scrcpy and self._scrcpy.running:
            return self._scrcpy.video_track

        logger.info("Preparing video track (starting scrcpy)...")
        info = await self._runtime.get_info()
        serial = info.metadata.get("serial", "")
        if not serial:
            logger.warning("No ADB serial found, cannot start scrcpy")
            return None

        self._scrcpy = ScrcpyAdapter(adb_serial=serial)
        video_track = await self._scrcpy.start()
        if self._webrtc_peer and video_track:
            self._webrtc_peer.set_video_track(video_track)
        if self._stream_lifecycle:
            await self._stream_lifecycle.viewer_connected()
        return video_track

    async def _on_start_stream(self) -> None:
        """Called when transitioning to STREAMING state — starts scrcpy and injects video track."""
        logger.info("Starting screen stream...")
        if self._scrcpy and self._scrcpy.running:
            return
        await self._prepare_video_track()

    async def _on_stop_stream(self) -> None:
        """Called when transitioning to IDLE state."""
        logger.info("Stopping screen stream...")
        if self._scrcpy:
            await self._scrcpy.stop()
            self._scrcpy = None
        if self._webrtc_peer:
            await self._webrtc_peer.close()

    # -- DataChannel control --

    async def _on_datachannel_message(self, msg: dict[str, Any]) -> None:
        """Handle DataChannel control messages (touch, key, etc.)."""
        msg_type = msg.get("type", "")

        if msg_type == "touch":
            action = msg.get("action", "")
            x, y = msg.get("x", 0), msg.get("y", 0)
            pointer_id = msg.get("pointerId", 0)

            if action == "down":
                self._touch_state[pointer_id] = (x, y)
            elif action == "move":
                prev = self._touch_state.get(pointer_id)
                if prev:
                    px, py = prev
                    if abs(x - px) > 5 or abs(y - py) > 5:
                        await self._runtime.swipe(px, py, x, y, 100)
                        self._touch_state[pointer_id] = (x, y)
            elif action == "up":
                start = self._touch_state.pop(pointer_id, None)
                if start:
                    sx, sy = start
                    dist = ((x - sx) ** 2 + (y - sy) ** 2) ** 0.5
                    if dist < 15:
                        await self._runtime.tap(x, y)
                    else:
                        await self._runtime.swipe(sx, sy, x, y, 200)

        elif msg_type == "scroll":
            x, y = msg.get("x", 0), msg.get("y", 0)
            dy = msg.get("dy", 0)
            direction = 1 if dy < 0 else -1
            await self._runtime.swipe(x, y, x, y + direction * 200, 200)

        elif msg_type == "key":
            keycode = msg.get("keycode", "")
            if keycode:
                await self._runtime.press_key(keycode)

        elif msg_type == "text":
            text = msg.get("content", "")
            if text:
                await self._runtime.type_text(text)

        elif msg_type == "back":
            await self._runtime.press_key("KEYCODE_BACK")

        elif msg_type == "home":
            await self._runtime.press_key("KEYCODE_HOME")

        elif msg_type == "configure":
            if self._scrcpy:
                self._scrcpy.configure(
                    max_fps=msg.get("maxFps"),
                    max_width=msg.get("maxWidth"),
                    bitrate=msg.get("bitrate"),
                )

        else:
            logger.debug("Unknown DataChannel message type: %s", msg_type)
