from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from cloud.internal_api.models import DeviceSnapshotModel, HeartbeatResponse
from core.errors import DeviceRuntimeError
from device.manager import DeviceManager
from host_agent.client import HostAgentClient
from host_agent.config import HostAgentConfig
from host_agent.policy_cache import HostPolicyCacheStore
from host_agent.status import AgentStatusTracker

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from host_agent.mcp_lock import McpBusyTracker


def build_device_snapshot(manager: DeviceManager) -> list[DeviceSnapshotModel]:
    return [
        DeviceSnapshotModel(
            device_id=device.id,
            driver_type=device.driver_type,
            status="idle" if device.status == "busy" else device.status,
            capability_tags=list(device.capability_tags),
        )
        for device in sorted(manager.list_devices(), key=lambda item: item.id)
    ]


class HeartbeatSynchronizer:
    def __init__(
        self,
        manager: DeviceManager,
        client: HostAgentClient,
        config: HostAgentConfig,
        *,
        address: str | None = None,
        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
        status_tracker: AgentStatusTracker | None = None,
        on_sync: Callable[[int], None] | None = None,
        policy_cache: HostPolicyCacheStore | None = None,
        on_policy_sync: Callable[[int], None] | None = None,
        mcp_busy_tracker: McpBusyTracker | None = None,
    ) -> None:
        self.manager = manager
        self.client = client
        self.config = config
        self.address = address
        self._sleep = sleep
        self.status_tracker = status_tracker
        self.on_sync = on_sync
        self.policy_cache = policy_cache
        self.on_policy_sync = on_policy_sync
        self.mcp_busy_tracker = mcp_busy_tracker
        self.policy = policy_cache.load() if policy_cache is not None else None
        self.policy_revision = self.policy.revision if self.policy is not None else 0
        if self.status_tracker is not None:
            self.status_tracker.mark_host_policy(self.policy)

    async def sync_once(self) -> HeartbeatResponse:
        self.probe_connected_devices()
        self.connect_devices()
        snapshot = build_device_snapshot(self.manager)
        mcp_busy_ids = (
            self.mcp_busy_tracker.busy_device_ids()
            if self.mcp_busy_tracker is not None
            else []
        )
        response = await self.client.heartbeat(
            snapshot,
            address=self.address,
            policy_revision=self.policy_revision,
            mcp_busy_device_ids=mcp_busy_ids,
        )
        self.policy_revision = response.policy_revision
        if response.policy is not None:
            self.policy = response.policy
            if self.policy_cache is not None:
                self.policy_cache.save(response.policy)
            if self.on_policy_sync is not None:
                self.on_policy_sync(response.policy.revision)
        elif response.policy_revision == 0:
            self.policy = None
            if self.policy_cache is not None:
                self.policy_cache.clear()
        if self.status_tracker is not None:
            self.status_tracker.mark_host_policy(self.policy)
        if self.status_tracker is not None:
            self.status_tracker.mark_heartbeat(ok=True, device_count=len(snapshot))
        if self.on_sync is not None:
            try:
                self.on_sync(len(snapshot))
            except Exception:
                pass
        return response

    async def run(self, stop: asyncio.Event) -> None:
        self.connect_devices()
        while not stop.is_set():
            await self.sync_once()
            try:
                await asyncio.wait_for(
                    stop.wait(),
                    timeout=self.config.heartbeat_interval_seconds,
                )
            except TimeoutError:
                continue

    def connect_devices(self) -> None:
        for device in self.manager.list_devices():
            if device.status not in {"idle", "offline", "error"}:
                continue
            try:
                self.manager.connect(device.id)
            except DeviceRuntimeError:
                continue

    def probe_connected_devices(self) -> None:
        active_id: str | None = None
        if self.status_tracker is not None:
            current = self.status_tracker.snapshot().get("current_assignment")
            if isinstance(current, dict) and isinstance(current.get("device_id"), str):
                active_id = current["device_id"]
        for device in self.manager.list_devices():
            if device.id != active_id and device.status == "busy":
                self.manager.probe(device.id)