"""Async Galaxy Repository client wrapper.

Wraps the read-only ``GalaxyRepository`` gRPC service exposed by the
MxAccess Gateway. The service lets callers test connectivity to the AVEVA
System Platform Galaxy Repository (ZB SQL database), read the last
deployment timestamp, and enumerate the deployed object hierarchy plus the
attributes on each object.
"""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator, Sequence
from datetime import datetime
from typing import Any

import grpc
from google.protobuf.timestamp_pb2 import Timestamp

from .auth import merge_metadata
from .errors import MxGatewayError, map_rpc_error
from .generated import galaxy_repository_pb2 as galaxy_pb
from .generated import galaxy_repository_pb2_grpc as galaxy_pb_grpc
from .options import (
    BrowseChildrenOptions,
    ClientOptions,
    create_channel,
    resolve_channel_security,
)

_DISCOVER_HIERARCHY_PAGE_SIZE = 5000
_BROWSE_CHILDREN_PAGE_SIZE = 500


class GalaxyRepositoryClient:
    """Async client for the Galaxy Repository gRPC service."""

    def __init__(
        self,
        *,
        options: ClientOptions,
        stub: Any,
        channel: grpc.aio.Channel | None = None,
    ) -> None:
        """Initialize the client with resolved options and a gRPC stub."""
        self.options = options
        self.raw_stub = stub
        self._channel = channel
        self._closed = False

    @classmethod
    async def connect(
        cls,
        options: ClientOptions | None = None,
        *,
        endpoint: str | None = None,
        api_key: str | None = None,
        plaintext: bool = False,
        ca_file: str | None = None,
        require_certificate_validation: bool = False,
        server_name_override: str | None = None,
        stub: Any | None = None,
    ) -> "GalaxyRepositoryClient":
        """Create a client with either a real async channel or a supplied fake stub."""

        resolved = options or ClientOptions(
            endpoint=endpoint or "",
            api_key=api_key,
            plaintext=plaintext,
            ca_file=ca_file,
            require_certificate_validation=require_certificate_validation,
            server_name_override=server_name_override,
        )

        if stub is not None:
            return cls(options=resolved, stub=stub)

        # Resolving security may perform a blocking TLS certificate probe (TOFU
        # default); run that off the event loop so connect never freezes it. The
        # channel itself must be built on the loop thread — a grpc.aio channel
        # binds to the loop current on the constructing thread, and a worker
        # thread has none.
        security = await asyncio.to_thread(resolve_channel_security, resolved)
        channel = create_channel(resolved, security=security)
        return cls(
            options=resolved,
            stub=galaxy_pb_grpc.GalaxyRepositoryStub(channel),
            channel=channel,
        )

    async def __aenter__(self) -> "GalaxyRepositoryClient":
        """Return self to support ``async with`` usage."""
        return self

    async def __aexit__(self, *_exc_info: object) -> None:
        """Close the client when leaving an ``async with`` block."""
        await self.close()

    async def close(self) -> None:
        """Close the owned gRPC channel."""

        if self._closed:
            return

        if self._channel is not None:
            await self._channel.close()
        self._closed = True

    async def test_connection(self) -> bool:
        """Return ``True`` when the gateway can reach the Galaxy Repository DB."""

        reply = await self._unary(
            "test connection",
            self.raw_stub.TestConnection,
            galaxy_pb.TestConnectionRequest(),
        )
        return bool(reply.ok)

    async def get_last_deploy_time(self) -> datetime | None:
        """Return the last Galaxy deploy timestamp or ``None`` when unset."""

        reply = await self._unary(
            "get last deploy time",
            self.raw_stub.GetLastDeployTime,
            galaxy_pb.GetLastDeployTimeRequest(),
        )
        if not reply.present:
            return None
        return reply.time_of_last_deploy.ToDatetime()

    async def discover_hierarchy(self) -> list[galaxy_pb.GalaxyObject]:
        """Return the deployed Galaxy object hierarchy as raw proto messages."""

        objects: list[galaxy_pb.GalaxyObject] = []
        seen_page_tokens: set[str] = set()
        page_token = ""
        while True:
            reply = await self._unary(
                "discover hierarchy",
                self.raw_stub.DiscoverHierarchy,
                galaxy_pb.DiscoverHierarchyRequest(
                    page_size=_DISCOVER_HIERARCHY_PAGE_SIZE,
                    page_token=page_token,
                ),
            )
            objects.extend(reply.objects)
            page_token = reply.next_page_token
            if not page_token:
                return objects
            if page_token in seen_page_tokens:
                raise MxGatewayError(
                    f"galaxy discover hierarchy returned repeated page token {page_token!r}"
                )
            seen_page_tokens.add(page_token)

    async def browse_children_raw(
        self, request: galaxy_pb.BrowseChildrenRequest
    ) -> galaxy_pb.BrowseChildrenReply:
        """Issue one BrowseChildren RPC and return the raw reply.

        Lower-level escape hatch for callers that need direct page-token control
        or do not want LazyBrowseNode wrapping. Most callers should use
        :py:meth:`browse` and :py:meth:`LazyBrowseNode.expand` instead.
        """

        return await self._unary(
            "browse children",
            self.raw_stub.BrowseChildren,
            request,
        )

    async def browse(
        self,
        options: BrowseChildrenOptions | None = None,
    ) -> list["LazyBrowseNode"]:
        """Return the root browse nodes for lazy hierarchy traversal.

        Each returned ``LazyBrowseNode`` wraps a Galaxy object whose direct
        children can be loaded on demand by ``await node.expand()``.
        """

        effective = options or BrowseChildrenOptions()
        return [
            node
            async for node in self._iter_browse_children(
                parent_gobject_id=None,
                options=effective,
            )
        ]

    async def _iter_browse_children(
        self,
        *,
        parent_gobject_id: int | None,
        options: BrowseChildrenOptions,
    ) -> AsyncIterator["LazyBrowseNode"]:
        page_token = ""
        seen_page_tokens: set[str] = set()
        while True:
            request = galaxy_pb.BrowseChildrenRequest(
                page_size=_BROWSE_CHILDREN_PAGE_SIZE,
                page_token=page_token,
                alarm_bearing_only=options.alarm_bearing_only,
                historized_only=options.historized_only,
            )
            if parent_gobject_id is not None:
                request.parent_gobject_id = parent_gobject_id
            if options.category_ids:
                request.category_ids.extend(options.category_ids)
            if options.template_chain_contains:
                request.template_chain_contains.extend(options.template_chain_contains)
            if options.tag_name_glob:
                request.tag_name_glob = options.tag_name_glob
            if options.include_attributes is not None:
                request.include_attributes = options.include_attributes

            reply = await self._unary(
                "browse children",
                self.raw_stub.BrowseChildren,
                request,
            )

            for index, obj in enumerate(reply.children):
                hint = (
                    index < len(reply.child_has_children)
                    and bool(reply.child_has_children[index])
                )
                yield LazyBrowseNode(self, obj, hint, options)

            page_token = reply.next_page_token
            if not page_token:
                return
            if page_token in seen_page_tokens:
                raise MxGatewayError(
                    f"galaxy browse children returned repeated page token {page_token!r}"
                )
            seen_page_tokens.add(page_token)

    def watch_deploy_events(
        self,
        last_seen_deploy_time: datetime | None = None,
        *,
        metadata: Sequence[tuple[str, str]] | None = None,
    ) -> AsyncIterator[galaxy_pb.DeployEvent]:
        """Stream Galaxy deploy events.

        On subscribe the gateway emits the current cached state and then one
        event per new deploy time. ``sequence`` is monotonic per server start;
        gaps mean events were dropped from the per-subscriber buffer. When
        ``last_seen_deploy_time`` is supplied and matches the current cached
        deploy time the bootstrap event is suppressed.
        """

        request = galaxy_pb.WatchDeployEventsRequest()
        if last_seen_deploy_time is not None:
            timestamp = Timestamp()
            timestamp.FromDatetime(last_seen_deploy_time)
            request.last_seen_deploy_time.CopyFrom(timestamp)

        kwargs: dict[str, Any] = {"metadata": merge_metadata(self.options.api_key, metadata)}
        if self.options.stream_timeout is not None:
            kwargs["timeout"] = self.options.stream_timeout
        try:
            call = self.raw_stub.WatchDeployEvents(request, **kwargs)
        except TypeError as error:
            if "timeout" not in kwargs or "unexpected keyword argument 'timeout'" not in str(error):
                raise
            kwargs.pop("timeout")
            call = self.raw_stub.WatchDeployEvents(request, **kwargs)

        return _canceling_iterator(call)

    async def _unary(
        self,
        operation: str,
        method: Any,
        request: Any,
        *,
        metadata: Sequence[tuple[str, str]] | None = None,
    ) -> Any:
        kwargs: dict[str, Any] = {"metadata": merge_metadata(self.options.api_key, metadata)}
        if self.options.call_timeout is not None:
            kwargs["timeout"] = self.options.call_timeout
        try:
            call = method(request, **kwargs)
        except TypeError as error:
            if "timeout" not in kwargs or "unexpected keyword argument 'timeout'" not in str(error):
                raise
            kwargs.pop("timeout")
            call = method(request, **kwargs)
        try:
            return await call
        except asyncio.CancelledError:
            cancel = getattr(call, "cancel", None)
            if cancel is not None:
                cancel()
            raise
        except grpc.RpcError as error:
            raise map_rpc_error(operation, error) from error


class LazyBrowseNode:
    """One node in a lazy-loaded Galaxy browse tree.

    Calling ``expand`` once fetches direct children (paginating as needed)
    and populates ``children``. Subsequent calls are no-ops so callers can
    drive UI expand toggles without de-duping.
    """

    def __init__(
        self,
        client: "GalaxyRepositoryClient",
        obj: galaxy_pb.GalaxyObject,
        has_children_hint: bool,
        options: BrowseChildrenOptions,
    ) -> None:
        """Initialize a node bound to its owning client and filter set."""
        self._client = client
        self._object = obj
        self._has_children_hint = has_children_hint
        self._options = options
        self._children: list[LazyBrowseNode] = []
        self._is_expanded = False
        self._expand_lock = asyncio.Lock()

    @property
    def object(self) -> galaxy_pb.GalaxyObject:
        """Return the underlying ``GalaxyObject`` proto for this node."""
        return self._object

    @property
    def has_children_hint(self) -> bool:
        """Return the server hint about whether this node has children."""
        return self._has_children_hint

    @property
    def children(self) -> list["LazyBrowseNode"]:
        """Return a copy of the loaded child nodes (empty until expanded)."""
        return list(self._children)

    @property
    def is_expanded(self) -> bool:
        """Return whether ``expand`` has already populated ``children``."""
        return self._is_expanded

    async def expand(self) -> None:
        """Fetch direct children of this node; no-op on subsequent calls."""
        if self._is_expanded:
            return
        async with self._expand_lock:
            if self._is_expanded:
                return
            new_children: list[LazyBrowseNode] = []
            async for child in self._client._iter_browse_children(
                parent_gobject_id=self._object.gobject_id,
                options=self._options,
            ):
                new_children.append(child)
            self._children.extend(new_children)
            self._is_expanded = True


async def _canceling_iterator(call: Any) -> AsyncIterator[galaxy_pb.DeployEvent]:
    try:
        async for event in call:
            yield event
    except asyncio.CancelledError:
        cancel = getattr(call, "cancel", None)
        if cancel is not None:
            cancel()
        raise
    except grpc.RpcError as error:
        raise map_rpc_error("watch deploy events", error) from error
    finally:
        cancel = getattr(call, "cancel", None)
        if cancel is not None:
            cancel()
