from __future__ import annotations

import json
import os
from dataclasses import dataclass, field, replace
from pathlib import Path
from secrets import token_urlsafe
from uuid import uuid4


class HostIdentityStateError(RuntimeError):
    """Raised when persisted Host enrollment identity is missing or invalid."""


@dataclass(frozen=True)
class HostIdentityState:
    agent_instance_id: str
    token: str = field(repr=False)
    host_id: str | None = None


class HostIdentityStore:
    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)

    def load(self) -> HostIdentityState | None:
        if not self.path.exists():
            return None
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
        except (OSError, ValueError, json.JSONDecodeError) as exc:
            raise HostIdentityStateError("Host identity state is unreadable") from exc
        if not isinstance(payload, dict):
            raise HostIdentityStateError("Host identity state must be an object")
        agent_instance_id = payload.get("agent_instance_id")
        token = payload.get("token")
        host_id = payload.get("host_id")
        if (
            not isinstance(agent_instance_id, str)
            or not agent_instance_id
            or not isinstance(token, str)
            or len(token) < 32
            or (host_id is not None and (not isinstance(host_id, str) or not host_id))
        ):
            raise HostIdentityStateError("Host identity state is invalid")
        return HostIdentityState(
            agent_instance_id=agent_instance_id,
            token=token,
            host_id=host_id,
        )

    def load_or_create(self) -> HostIdentityState:
        existing = self.load()
        if existing is not None:
            return existing
        state = HostIdentityState(
            agent_instance_id=f"agent-{uuid4().hex}",
            token=token_urlsafe(48),
        )
        self._write(state)
        return state

    def complete(self, state: HostIdentityState, host_id: str) -> HostIdentityState:
        if not host_id:
            raise ValueError("host_id must not be empty")
        current = self.load()
        if current is not None and (
            current.agent_instance_id != state.agent_instance_id
            or current.token != state.token
        ):
            raise HostIdentityStateError("Host identity changed during enrollment")
        completed = replace(state, host_id=host_id)
        self._write(completed)
        return completed

    def _write(self, state: HostIdentityState) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp")
        payload = {
            "agent_instance_id": state.agent_instance_id,
            "token": state.token,
            "host_id": state.host_id,
        }
        try:
            temporary.write_text(
                json.dumps(payload, ensure_ascii=True, indent=2) + "\n",
                encoding="utf-8",
            )
            _restrict_permissions(temporary)
            os.replace(temporary, self.path)
            _restrict_permissions(self.path)
        finally:
            if temporary.exists():
                temporary.unlink()


def _restrict_permissions(path: Path) -> None:
    try:
        path.chmod(0o600)
    except OSError:
        return