# ruff: noqa: D102
from __future__ import annotations

import asyncio
import dataclasses
import functools
import logging
import sys
import threading
import time
import uuid
from collections.abc import AsyncGenerator
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from datetime import timedelta
from typing import Any
from typing import ClassVar
from typing import Generic
from typing import Literal
from typing import NamedTuple
from typing import TYPE_CHECKING

from academy.exchange.cloud.client import DEFAULT_EXCHANGE_URL

if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
    from typing import Self
else:  # pragma: <3.11 cover
    from typing_extensions import Self


import globus_sdk
from globus_sdk import AuthClient
from globus_sdk import DependentScopeSpec
from globus_sdk import GlobusApp
from globus_sdk import GlobusHTTPResponse
from globus_sdk import Scope
from globus_sdk.authorizers import GlobusAuthorizer
from globus_sdk.exc import GlobusAPIError
from globus_sdk.gare import GlobusAuthorizationParameters
from globus_sdk.scopes import AuthScopes
from globus_sdk.transport.requests import RequestsTransport
from pydantic import BaseModel
from pydantic import Field

from academy.exception import BadEntityIdError
from academy.exception import MailboxTerminatedError
from academy.exception import UnauthorizedError
from academy.exchange.cloud.app import StatusCode
from academy.exchange.cloud.login import get_globus_app
from academy.exchange.cloud.scopes import AcademyExchangeScopes
from academy.exchange.cloud.scopes import get_academy_exchange_scope_id
from academy.exchange.factory import ExchangeFactory
from academy.exchange.transport import ExchangeTransportMixin
from academy.exchange.transport import MailboxStatus
from academy.identifier import AgentId
from academy.identifier import EntityId
from academy.identifier import UserId
from academy.message import Message
from academy.serialize import NoPickleMixin

if TYPE_CHECKING:
    from academy.agent import Agent
    from academy.agent import AgentT
else:
    from academy.identifier import AgentT

logger = logging.getLogger(__name__)


class _AcademyConnectionInfo(NamedTuple):
    project_id: uuid.UUID
    client_params: dict[str, Any] | None = None
    request_timeout_s: float = 60


class AcademyAPIError(GlobusAPIError):
    """Error class to represent error responses from Academy."""


class AcademyGlobusClient(globus_sdk.BaseClient):
    """A globus service client to make requests to hosted exchange.

    The GlobusExchangeClient acts as a wrapper through which authenticated
    requests are issued. The exchange automatically handles things like
    retrying, refreshing tokens, and exponential backoff.
    The [`BaseClient`][globus_sdk.BaseClient] is implemented using requests,
    so calls are synchronous, and must be run within a thread when using
    asyncio.
    """

    base_url = DEFAULT_EXCHANGE_URL
    scopes = AcademyExchangeScopes
    default_scope_requirements: ClassVar[list[Scope]] = [
        AcademyExchangeScopes.academy_exchange,
    ]
    error_class = AcademyAPIError

    _mailbox_url = '/mailbox'
    _message_url = '/message'
    _discover_url = '/discover'
    _heartbeat_url = '/mailbox/heartbeat'

    def discover(
        self,
        agent: type[Agent] | str,
        *,
        allow_subclasses: bool = True,
    ) -> GlobusHTTPResponse:
        if isinstance(agent, str):
            agent_str = agent
        else:
            agent_str = f'{agent.__module__}.{agent.__name__}'
        return self.request(
            'GET',
            self._discover_url,
            data={
                'agent': agent_str,
                'allow_subclasses': allow_subclasses,
            },
        )

    def recv(
        self,
        mailbox_id: EntityId,
        timeout: float,
    ) -> GlobusHTTPResponse:
        return self.request(
            'GET',
            self._message_url,
            data={
                'mailbox': mailbox_id.model_dump_json(),
                'timeout': timeout,
            },
        )

    def register_agent(
        self,
        agent_id: AgentId[AgentT],
        agent: type[AgentT],
    ) -> GlobusHTTPResponse:
        return self.post(
            self._mailbox_url,
            data={
                'mailbox': agent_id.model_dump_json(),
                'agent': ','.join(agent._agent_mro()),
            },
        )

    def register_client(
        self,
        mailbox_id: EntityId,
    ) -> GlobusHTTPResponse:
        return self.post(
            self._mailbox_url,
            data={
                'mailbox': mailbox_id.model_dump_json(),
            },
        )

    def send(self, message: Message[Any]) -> GlobusHTTPResponse:
        return self.put(
            self._message_url,
            data={'message': message.model_dump_json()},
        )

    def status(self, uid: EntityId) -> GlobusHTTPResponse:
        return self.request(
            'GET',
            self._mailbox_url,
            data={'mailbox': uid.model_dump_json()},
        )

    def terminate(self, uid: EntityId) -> GlobusHTTPResponse:
        return self.request(
            'DELETE',
            self._mailbox_url,
            data={'mailbox': uid.model_dump_json()},
        )

    def get_heartbeat(self, uid: EntityId) -> GlobusHTTPResponse:
        return self.request(
            'GET',
            self._heartbeat_url,
            data={'mailbox': uid.model_dump_json()},
        )


@dataclasses.dataclass
class _PendingRegistration(Generic[AgentT]):
    """Intermediate state for a registration awaiting token delegation.

    Created during the prepare phase of batch registration. Contains
    the Globus Auth client, credential, and scope but not yet the
    delegated token (which requires user consent).
    """

    agent_id: AgentId[AgentT]
    client_id: uuid.UUID
    secret: str
    scope: Scope


class GlobusAgentRegistration(BaseModel, Generic[AgentT]):
    """Agent registration for hosted globus exchange."""

    agent_id: AgentId[AgentT]
    """Unique identifier for the agent created by the exchange."""

    client_id: uuid.UUID
    """Client ID of Globus resource server.

    Each agent is a resources server.  This allows the agent to exchange
    delegated tokens to act on the exchange on behalf of the client, and
    to create it's own delegated tokens so other agents can act on its
    behalf.
    """

    token: str
    """Auth. token provided by launching client (user or another agent)."""

    secret: str
    """Secret for agent to use to authenticate itself with Globus Auth.

    Agents are created as hybrid resource servers within the Globus ecosystem.
    This allows them to use delegated tokens, but also requires the them to
    be able to store secrets. In order to support the launching of agents,
    we pass a secret as part of the agent initialization. We assume the
    security of the launching channel (typically Globus Compute in the
    Academy ecosystem).
    """

    exchange_type: Literal['globus'] = Field('globus', repr=False)


class GlobusExchangeTransport(ExchangeTransportMixin, NoPickleMixin):
    """Globus exchange client.

    Args:
        mailbox_id: Identifier of the mailbox on the exchange. If there is
            not an id provided, the exchange will create a new client mailbox.
        connection_info: Project id, client parameters and other information
            about the connection to the service.
        app: For user authorization through token retrieval.
        authorizer: For service authorization through token retrieval.
    """

    def __init__(
        self,
        mailbox_id: EntityId,
        *,
        connection_info: _AcademyConnectionInfo,
        app: GlobusApp | None = None,
        authorizer: GlobusAuthorizer | None = None,
    ) -> None:
        self._mailbox_id = mailbox_id
        self.project = connection_info.project_id
        self.child_clients: list[uuid.UUID] = []
        self.client_params = connection_info.client_params or {}
        self.request_timeout_s = connection_info.request_timeout_s

        self.login_time = datetime.min
        self._app = app
        self._authorizer = authorizer
        self._auth_lock = threading.Lock()
        self._local_data = threading.local()
        self.executor = ThreadPoolExecutor(
            thread_name_prefix='exchange-globus-thread',
        )

    @property
    def exchange_client(self) -> AcademyGlobusClient:
        """A thread local copy of the Globus AuthClient."""
        try:
            return self._local_data.exchange_client
        except AttributeError:
            self._local_data.exchange_client = AcademyGlobusClient(
                app=self._app,
                authorizer=self._authorizer,
                transport=RequestsTransport(http_timeout=-1),
                **self.client_params,
            )
            return self._local_data.exchange_client

    @property
    def auth_client(self) -> AuthClient:
        """A thread local copy of the Globus AuthClient."""
        # Fast path: return cached thread-local client.
        if datetime.now() - self.login_time < timedelta(minutes=30):
            try:
                return self._local_data.auth_client
            except AttributeError:
                logger.debug(
                    'Auth client does not exist for thread.',
                )
        else:
            # Login required — serialize to prevent duplicate
            # interactive prompts from concurrent threads.
            with self._auth_lock:
                if datetime.now() - self.login_time >= timedelta(minutes=30):
                    logger.debug('Auth client login timeout.')

                    if self._app is None:
                        raise NotImplementedError(
                            'Launching child agents is'
                            ' currently not implemented.',
                        )

                    self._app.add_scope_requirements(
                        {
                            AuthScopes.resource_server: [
                                AuthScopes.manage_projects,
                                AuthScopes.email,
                                AuthScopes.profile,
                            ],
                        },
                    )

                    self._app.login(
                        force=True,
                        auth_params=GlobusAuthorizationParameters(
                            prompt='login',
                        ),
                    )
                    self.login_time = datetime.now()

        logger.info('Initializing auth client.')
        self._local_data.auth_client = AuthClient(
            app=self._app,
        )
        return self._local_data.auth_client

    def _register_client(self) -> None:
        self.exchange_client.register_client(self.mailbox_id)

    @classmethod
    async def new(
        cls,
        *,
        connection_info: _AcademyConnectionInfo,
        app: GlobusApp | None = None,
        authorizer: GlobusAuthorizer | None = None,
        mailbox_id: EntityId | None = None,
        name: str | None = None,
    ) -> Self:
        """Instantiate a new transport.

        Args:
            connection_info: Project id, client parameters and other
                information about the connection to the service.
            app: For user authorization through token retrieval
            authorizer: For service authorization through token retrieval
            mailbox_id: Bind the transport to the specific mailbox. If `None`,
                a new user entity will be registered and the transport will be
                bound to that mailbox.
            name: Display name of the registered entity if `mailbox_id` is
                `None`.

        Returns:
            An instantiated transport bound to a specific mailbox.
        """
        loop = asyncio.get_running_loop()

        if mailbox_id is None:
            mailbox_id = UserId.new(name=name)
            client = cls(
                mailbox_id,
                connection_info=connection_info,
                app=app,
                authorizer=authorizer,
            )
            await loop.run_in_executor(
                client.executor,
                client._register_client,
            )
            logger.info(
                'Registered %s in exchange',
                mailbox_id,
                extra={'academy.mailbox_id': mailbox_id},
            )
            return client

        return cls(
            mailbox_id,
            connection_info=connection_info,
            app=app,
            authorizer=authorizer,
        )

    @property
    def mailbox_id(self) -> EntityId:
        return self._mailbox_id

    def _clean_up(self) -> None:
        for client_id in self.child_clients:
            self.auth_client.delete_client(client_id)

    async def close(self) -> None:
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(self.executor, self._clean_up)

    def _discover(
        self,
        agent: type[Agent] | str,
        allow_subclasses: bool,
    ) -> GlobusHTTPResponse:
        return self.exchange_client.discover(
            agent,
            allow_subclasses=allow_subclasses,
        )

    async def discover(
        self,
        agent: type[Agent] | str,
        *,
        allow_subclasses: bool = True,
    ) -> tuple[AgentId[Any], ...]:
        loop = asyncio.get_running_loop()
        response = await loop.run_in_executor(
            self.executor,
            self._discover,
            agent,
            allow_subclasses,
        )

        agent_ids_str = response['agent_ids']
        agent_ids = [aid for aid in agent_ids_str.split(',') if len(aid) > 0]
        return tuple(AgentId(uid=uuid.UUID(aid)) for aid in agent_ids)

    def factory(self) -> GlobusExchangeFactory:
        return GlobusExchangeFactory(self.project, self.client_params)

    def _recv_sync(self, timeout: float) -> GlobusHTTPResponse:
        return self.exchange_client.recv(self.mailbox_id, timeout)

    async def _recv(self, timeout: float | None = None) -> Message[Any]:
        start_time = time.time()
        current_time = time.time()
        loop = asyncio.get_running_loop()
        while timeout is None or current_time - start_time < timeout:
            internal_timeout = (
                self.request_timeout_s
                if timeout is None
                else min(
                    (start_time + timeout) - current_time,
                    self.request_timeout_s,
                )
            )
            try:
                logger.debug(
                    f'Receiving message for mailbox {self.mailbox_id}',
                    extra={'academy.mailbox_id': self.mailbox_id},
                )
                response = await asyncio.wait_for(
                    loop.run_in_executor(
                        self.executor,
                        self._recv_sync,
                        internal_timeout,
                    ),
                    timeout,
                )
                message_raw = response['message']
                logger.debug(
                    f'Received message of length {len(message_raw)}',
                    extra={'academy.message_length': len(message_raw)},
                )
                break
            except AcademyAPIError as e:
                if e.http_status == StatusCode.TERMINATED.value:
                    raise MailboxTerminatedError(self.mailbox_id) from e
                elif (
                    e.http_status == StatusCode.TIMEOUT.value
                ):  # pragma: no cover
                    logger.debug(
                        f'Failed to receive response in {internal_timeout} '
                        'seconds.',
                    )
                    pass
            except asyncio.TimeoutError:
                logger.debug(
                    f'Failed to receive response in {internal_timeout} '
                    'seconds.',
                )
                pass
            current_time = time.time()
        else:
            raise TimeoutError(
                f'Failed to receive response in {timeout} seconds.',
            )

        return Message.model_validate_json(message_raw)

    async def listen(
        self,
        timeout: float | None = None,
    ) -> AsyncGenerator[Message[Any]]:
        while True:
            yield await self._recv(timeout)

    def _prepare_registration(
        self,
        aid: AgentId[AgentT],
    ) -> _PendingRegistration[AgentT]:
        """Create Globus Auth client, credential, and scope for an agent.

        This is phase 1 of registration. It creates all Globus Auth
        entities but does **not** accumulate scope requirements or
        trigger user consent. The caller is responsible for adding
        the returned scope via
        ``GlobusApp.add_scope_requirements`` and triggering login
        before calling ``_finalize_registration``.
        """
        # Create new resource server (auth entity)
        logger.info('Creating new auth client')
        client_response = self.auth_client.create_client(
            str(aid),
            project=self.project,
            client_type='hybrid_confidential_client_resource_server',
            visibility='private',
        )
        client_id = client_response['client']['id']
        logger.info(
            f'Created client id: {client_id}',
            extra={'academy.client_id': client_id},
        )

        # Create secret
        logger.info('Creating secret.')
        credentials_response = self.auth_client.create_client_credential(
            client_id,
            'Launch Credentials',
        )
        secret = credentials_response['credential']['secret']

        # Create scope
        logger.info('Creating new scope.')
        scope_response = self.auth_client.create_scope(
            client_id,
            'Agent launch',
            'Launch agent',
            'launch',
            dependent_scopes=[
                DependentScopeSpec(
                    get_academy_exchange_scope_id(),
                    optional=False,
                    requires_refresh_token=True,
                ),
            ],
            allows_refresh_token=True,
        )
        scope = Scope.parse(scope_response['scopes'][0]['scope_string'])

        return _PendingRegistration(
            agent_id=aid,
            client_id=client_id,
            secret=secret,
            scope=scope,
        )

    def _finalize_registration(
        self,
        pending: _PendingRegistration[AgentT],
    ) -> GlobusAgentRegistration[AgentT]:
        """Obtain delegated token for a previously prepared agent.

        This is phase 2 of registration. The GlobusApp must already
        have valid tokens for the agent's scope (via a prior
        ``login()`` call that covered all accumulated scopes).
        """
        logger.info('Creating delegated token.')
        assert self._app is not None
        authorizer = self._app.get_authorizer(str(pending.client_id))
        bearer = authorizer.get_authorization_header()
        if bearer is None:  # pragma: no cover
            raise UnauthorizedError(
                'Unable to get authorization headers.',
            )

        auth_header_parts = bearer.split(' ')
        token = auth_header_parts[1]

        return GlobusAgentRegistration(
            agent_id=pending.agent_id,
            client_id=pending.client_id,
            token=token,
            secret=pending.secret,
        )

    def _create_registration(
        self,
        aid: AgentId[AgentT],
    ) -> GlobusAgentRegistration[AgentT]:
        """Create a single agent registration (prepare + finalize).

        Convenience method that runs both phases sequentially for a
        single agent. Triggers a consent prompt for the new scope.
        """
        pending = self._prepare_registration(aid)
        assert self._app is not None
        self._app.add_scope_requirements(
            {str(pending.client_id): [pending.scope]},
        )
        return self._finalize_registration(pending)

    def _create_mailbox(
        self,
        aid: AgentId[AgentT],
        agent: type[AgentT],
    ) -> None:
        self.exchange_client.register_agent(aid, agent)

    async def register_agent(
        self,
        agent: type[AgentT],
        *,
        name: str | None = None,
    ) -> GlobusAgentRegistration[AgentT]:
        loop = asyncio.get_running_loop()

        # Create registration
        aid: AgentId[AgentT] = AgentId.new(name=name)
        registration = await loop.run_in_executor(
            self.executor,
            self._create_registration,
            aid,
        )

        # Create mailbox
        await loop.run_in_executor(
            self.executor,
            self._create_mailbox,
            aid,
            agent,
        )

        self.child_clients.append(registration.client_id)

        return registration

    async def register_agents(
        self,
        agents: list[tuple[type[AgentT], str | None]],
    ) -> list[GlobusAgentRegistration[AgentT]]:
        """Register multiple agents with a single auth prompt.

        Args:
            agents: List of (agent_type, name) pairs to register.

        Returns:
            List of agent registrations in order of input.
        """
        loop = asyncio.get_running_loop()

        # Phase 1: create auth entities concurrently.
        # _prepare_registration no longer mutates shared GlobusApp
        # state, and the auth_client property serializes interactive
        # login via _auth_lock, so parallel dispatch is safe.
        aids: list[AgentId[AgentT]] = [
            AgentId.new(name=name) for _, name in agents
        ]
        prepared: list[_PendingRegistration[AgentT]] = list(
            await asyncio.gather(
                *(
                    loop.run_in_executor(
                        self.executor,
                        self._prepare_registration,
                        aid,
                    )
                    for aid in aids
                ),
            ),
        )
        pending = list(
            zip(
                prepared,
                [at for at, _ in agents],
                strict=True,
            ),
        )

        # Accumulate scope requirements, then do a single login
        # that satisfies all of them at once.
        assert self._app is not None
        for p, _ in pending:
            self._app.add_scope_requirements(
                {str(p.client_id): [p.scope]},
            )

        if self._app.login_required():
            await loop.run_in_executor(
                self.executor,
                self._app.login,
            )

        # Phase 2: extract delegated tokens and create mailboxes.
        # Safe to parallelize — no shared state mutation, and the
        # app already holds valid tokens for all scopes.
        regs: list[GlobusAgentRegistration[AgentT]] = list(
            await asyncio.gather(
                *(
                    loop.run_in_executor(
                        self.executor,
                        self._finalize_registration,
                        p,
                    )
                    for p, _ in pending
                ),
            ),
        )

        await asyncio.gather(
            *(
                loop.run_in_executor(
                    self.executor,
                    self._create_mailbox,
                    reg.agent_id,
                    agent_type,
                )
                for reg, (_, agent_type) in zip(
                    regs,
                    pending,
                    strict=True,
                )
            ),
        )

        for reg in regs:
            self.child_clients.append(reg.client_id)

        return regs

    def _send(self, message: Message[Any]) -> None:
        self.exchange_client.send(message)

    async def send(self, message: Message[Any]) -> None:
        loop = asyncio.get_running_loop()
        try:
            await loop.run_in_executor(
                self.executor,
                self._send,
                message,
            )
        except AcademyAPIError as e:
            if e.http_status == StatusCode.NOT_FOUND.value:
                raise BadEntityIdError(message.dest) from e
            elif e.http_status == StatusCode.TERMINATED.value:
                raise MailboxTerminatedError(message.dest) from e
            raise e  # pragma: no cover

    def _status(self, uid: EntityId) -> MailboxStatus:
        response = self.exchange_client.status(uid)
        status = response['status']
        return MailboxStatus(status)

    async def status(self, uid: EntityId) -> MailboxStatus:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            self.executor,
            self._status,
            uid,
        )

    def _terminate(self, uid: EntityId) -> None:
        self.exchange_client.terminate(uid)

    async def terminate(self, uid: EntityId) -> None:
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(
            self.executor,
            self._terminate,
            uid,
        )

    def _get_heartbeat(self, uid: EntityId) -> float | None:
        missing_code = 404
        try:
            response = self.exchange_client.get_heartbeat(uid)
            return response.get('heartbeat')
        except AcademyAPIError as e:
            if e.http_status == missing_code:
                raise BadEntityIdError(uid) from e
            raise

    async def heartbeat_status(self, uid: EntityId) -> float | None:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            self.executor,
            self._get_heartbeat,
            uid,
        )

    async def update_heartbeat(self) -> None:
        pass  # Server tracks this automatically via listen/send


class GlobusExchangeFactory(ExchangeFactory[GlobusExchangeTransport]):
    """Globus exchange client factory.

    Args:
        project_id: Project to create new clients under. Must be able
            to authenticate as a administrator.
        client_params: Additional parameters for globus client.
        request_timeout_s: Maximum length of receive/listen requests to
            the exchange.
    """

    def __init__(
        self,
        project_id: uuid.UUID,
        client_params: dict[str, Any] | None = None,
        request_timeout_s: float = 60,
    ) -> None:
        self.info = _AcademyConnectionInfo(
            project_id=project_id,
            client_params=client_params,
            request_timeout_s=request_timeout_s,
        )

    async def _create_transport(
        self,
        mailbox_id: EntityId | None = None,
        *,
        name: str | None = None,
        registration: GlobusAgentRegistration[Any] | None = None,  # type: ignore[override]
    ) -> GlobusExchangeTransport:
        loop = asyncio.get_running_loop()

        if registration is None:
            app = await loop.run_in_executor(
                None,
                get_globus_app,
            )
            return await GlobusExchangeTransport.new(
                app=app,
                mailbox_id=mailbox_id,
                name=name,
                connection_info=self.info,
            )
        else:
            logger.info('Initializing auth client for new agent.')
            auth_client = globus_sdk.ConfidentialAppAuthClient(
                client_id=registration.client_id,
                client_secret=registration.secret,
            )

            logger.info('Exchanging dependent token.')
            dependent_token_response = await loop.run_in_executor(
                None,
                functools.partial(
                    auth_client.oauth2_get_dependent_tokens,
                    registration.token,
                    refresh_tokens=True,
                    scope=str(AcademyExchangeScopes.academy_exchange),
                ),
            )

            # TODO: dependent token only provides access to exchange
            # So agents do not have access to auth. to create new agents.
            authorizer = globus_sdk.RefreshTokenAuthorizer(
                dependent_token_response.by_resource_server[
                    AcademyExchangeScopes.resource_server
                ]['refresh_token'],
                auth_client,
            )

            return await GlobusExchangeTransport.new(
                authorizer=authorizer,
                mailbox_id=mailbox_id,
                name=name,
                connection_info=self.info,
            )
