"""Agent registration and enrollment."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import structlog

from curadise_agent.errors import RegistrationError
from curadise_agent.state.manager import AgentState

if TYPE_CHECKING:
    from curadise_agent.state.manager import StateManager
    from curadise_agent.transport.client import TransportClient

log = structlog.get_logger(__name__)


@dataclass
class RegistrationResult:
    """Result of agent registration."""

    agent_id: str
    access_token: str
    refresh_token: str | None = None
    expires_in: int | None = None


class RegistrationService:
    """
    Handles agent registration with the Curadise server.

    Manages the enrollment token exchange process and stores
    resulting credentials.
    """

    def __init__(
        self,
        transport: TransportClient,
        state_manager: StateManager,
    ) -> None:
        """
        Initialize registration service.

        Args:
            transport: HTTP transport client
            state_manager: State manager for persistence
        """
        self._transport = transport
        self._state = state_manager

    async def register(
        self,
        enrollment_token: str,
        system_info: dict[str, Any],
        tags: dict[str, str] | None = None,
    ) -> RegistrationResult:
        """
        Register the agent with the server.

        Args:
            enrollment_token: One-time enrollment token
            system_info: System information to send
            tags: Optional agent tags

        Returns:
            RegistrationResult with agent ID and tokens

        Raises:
            RegistrationError: If registration fails
        """
        self._state.transition_to(AgentState.REGISTERING)

        try:
            log.info("registration_starting")

            response = await self._transport.register(
                enrollment_token=enrollment_token,
                system_info=system_info,
                tags=tags,
            )

            # Extract credentials from response
            agent_id = response.get("agent_id")
            access_token = response.get("access_token")
            refresh_token = response.get("refresh_token")
            expires_in = response.get("expires_in")

            if not agent_id:
                raise RegistrationError("Server response missing agent_id")
            if not access_token:
                raise RegistrationError("Server response missing access_token")

            # Store credentials
            self._state.set_agent_id(agent_id)
            self._state.set_token(access_token)

            # Update transport with new token
            self._transport.set_auth_token(access_token)

            # Store additional metadata
            if refresh_token:
                self._state.set_metadata("refresh_token", refresh_token)
            if expires_in:
                self._state.set_metadata("token_expires_in", expires_in)

            log.info("registration_successful", agent_id=agent_id)

            return RegistrationResult(
                agent_id=agent_id,
                access_token=access_token,
                refresh_token=refresh_token,
                expires_in=expires_in,
            )

        except RegistrationError:
            self._state.transition_to(AgentState.INITIALIZING)
            raise
        except Exception as e:
            self._state.transition_to(AgentState.INITIALIZING)
            log.error("registration_failed", error=str(e))
            raise RegistrationError(f"Registration failed: {e}", cause=e) from e

    async def check_registration(self) -> bool:
        """
        Check if agent has valid registration.

        Returns:
            True if registered with valid credentials
        """
        if not self._state.is_registered:
            return False

        if not self._state.is_authenticated:
            return False

        # Optionally verify with server
        try:
            # Could make a lightweight API call to verify
            return True
        except Exception:
            return False

    async def unregister(self) -> None:
        """
        Unregister the agent from the server.

        Clears local credentials and optionally notifies server.
        """
        if not self._state.is_registered:
            log.warning("unregister_not_registered")
            return

        agent_id = self._state.agent_id
        log.info("unregistering", agent_id=agent_id)

        # Clear transport auth
        self._transport.clear_auth_token()

        # Clear stored credentials
        self._state.clear_credentials()

        log.info("unregistered", agent_id=agent_id)
