from __future__ import annotations

from typing import Any

from device.manager import DeviceManager
from driver.registry import build_driver_factory
from host_agent.client import HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig
from storage.device_config import DeviceConfigStore


def register_local_device(
    config_store: DeviceConfigStore,
    manager: DeviceManager,
    *,
    device_id: str,
    driver_type: str,
    connection_info: dict[str, Any],
    name: str | None,
    config: HostAgentConfig,
    enrollment_client: HostAgentEnrollmentClient | None,
) -> None:
    previous = config_store.get(device_id)
    previous_runtime_id = (
        previous["cloud_device_id"] or previous["device_id"] if previous else None
    )
    config_store.add(
        device_id=device_id,
        name=name,
        driver_type=driver_type,
        connection_info=connection_info,
    )
    runtime_device_id = device_id
    if config.enrollment_managed:
        assert enrollment_client is not None
        enrollment = enrollment_client.enroll_device(
            local_device_id=device_id,
            driver_type=driver_type,
            name=name,
            capability_tags=[],
        )
        runtime_device_id = enrollment.device_id
        config_store.set_cloud_device_id(device_id, runtime_device_id)
    if previous_runtime_id is not None and previous_runtime_id != runtime_device_id:
        manager.unregister_device(previous_runtime_id)
    manager.register_device(
        runtime_device_id,
        build_driver_factory(driver_type, connection_info),
        name=name,
        driver_type=driver_type,
        connection_info=connection_info,
    )


def unregister_local_device(
    config_store: DeviceConfigStore,
    manager: DeviceManager,
    *,
    device_id: str,
) -> None:
    record = config_store.get(device_id)
    if record is None:
        return
    runtime_device_id = record["cloud_device_id"] or record["device_id"]
    manager.unregister_device(runtime_device_id)
    config_store.remove(device_id)