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

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable


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,
    ) -> None:
        self.manager = manager
        self.client = client
        self.config = config
        self.address = address
        self._sleep = sleep

    async def sync_once(self) -> HeartbeatResponse:
        return await self.client.heartbeat(
            build_device_snapshot(self.manager),
            address=self.address,
        )

    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 != "idle":
                continue
            try:
                self.manager.connect(device.id)
            except DeviceRuntimeError:
                continue