from __future__ import annotations

from dataclasses import replace

from host_agent.client import HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig, HostAgentConfigurationError
from host_agent.identity import HostIdentityStore


def resolve_host_identity(
    config: HostAgentConfig,
    *,
    identity_store: HostIdentityStore,
    client: HostAgentEnrollmentClient,
) -> HostAgentConfig:
    if config.host_id and config.token:
        return config
    state = identity_store.load_or_create()
    if state.host_id is None:
        if not config.enrollment_token:
            raise HostAgentConfigurationError(
                "HOST_AGENT_ENROLLMENT_TOKEN is required to complete enrollment"
            )
        response = client.enroll_host(
            agent_instance_id=state.agent_instance_id,
            host_token=state.token,
            display_name=config.display_name,
        )
        state = identity_store.complete(state, response.host_id)
    return replace(
        config,
        host_id=state.host_id or "",
        token=state.token,
        enrollment_managed=True,
    )