"""Tests for the Galaxy Repository async client wrapper."""

from __future__ import annotations

import asyncio
from datetime import datetime, timezone
from typing import Any

import grpc
import pytest
from google.protobuf.timestamp_pb2 import Timestamp

from zb_mom_ww_mxgateway import ClientOptions, DeployEvent, GalaxyRepositoryClient, WatchDeployEventsRequest
from zb_mom_ww_mxgateway.errors import MxGatewayError
from zb_mom_ww_mxgateway.galaxy import LazyBrowseNode
from zb_mom_ww_mxgateway.generated import galaxy_repository_pb2 as galaxy_pb
from zb_mom_ww_mxgateway.generated import galaxy_repository_pb2_grpc as galaxy_pb_grpc
from zb_mom_ww_mxgateway.options import BrowseChildrenOptions


def test_galaxy_messages_import() -> None:
    request = galaxy_pb.DiscoverHierarchyRequest()
    obj = galaxy_pb.GalaxyObject(
        gobject_id=42,
        tag_name="DelmiaReceiver_001",
        contained_name="DelmiaReceiver",
        browse_name="DelmiaReceiver",
        parent_gobject_id=10,
        is_area=False,
        category_id=4,
        hosted_by_gobject_id=10,
        template_chain=["$ApplicationObject", "$DelmiaReceiver"],
        attributes=[
            galaxy_pb.GalaxyAttribute(
                attribute_name="DownloadPath",
                full_tag_reference="DelmiaReceiver_001.DownloadPath",
                mx_data_type=8,
                data_type_name="String",
            ),
        ],
    )

    assert request.DESCRIPTOR is not None
    assert obj.attributes[0].attribute_name == "DownloadPath"
    assert hasattr(galaxy_pb_grpc, "GalaxyRepositoryStub")


@pytest.mark.asyncio
async def test_test_connection_returns_bool_and_sends_auth() -> None:
    stub = FakeGalaxyStub()
    stub.test_connection.replies = [galaxy_pb.TestConnectionReply(ok=True)]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    result = await client.test_connection()

    assert result is True
    assert stub.test_connection.metadata == (("authorization", "Bearer mxgw_test_secret"),)
    assert isinstance(stub.test_connection.requests[0], galaxy_pb.TestConnectionRequest)


@pytest.mark.asyncio
async def test_get_last_deploy_time_returns_datetime_when_present() -> None:
    timestamp = Timestamp()
    timestamp.FromDatetime(datetime(2025, 4, 1, 12, 30, 45, tzinfo=timezone.utc))
    stub = FakeGalaxyStub()
    stub.get_last_deploy_time.replies = [
        galaxy_pb.GetLastDeployTimeReply(present=True, time_of_last_deploy=timestamp),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    when = await client.get_last_deploy_time()

    assert when is not None
    assert when.year == 2025
    assert when.month == 4
    assert when.day == 1
    assert when.hour == 12
    assert when.minute == 30
    assert when.second == 45


@pytest.mark.asyncio
async def test_get_last_deploy_time_returns_none_when_not_present() -> None:
    stub = FakeGalaxyStub()
    stub.get_last_deploy_time.replies = [galaxy_pb.GetLastDeployTimeReply(present=False)]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    assert await client.get_last_deploy_time() is None


@pytest.mark.asyncio
async def test_discover_hierarchy_returns_proto_objects() -> None:
    stub = FakeGalaxyStub()
    stub.discover_hierarchy.replies = [
        galaxy_pb.DiscoverHierarchyReply(
            next_page_token="page-2",
            total_object_count=2,
            objects=[
                galaxy_pb.GalaxyObject(
                    gobject_id=1,
                    tag_name="TestMachine_001",
                    contained_name="TestMachine",
                    browse_name="TestMachine_001",
                    is_area=True,
                ),
            ],
        ),
        galaxy_pb.DiscoverHierarchyReply(
            total_object_count=2,
            objects=[
                galaxy_pb.GalaxyObject(
                    gobject_id=2,
                    tag_name="DelmiaReceiver_001",
                    contained_name="DelmiaReceiver",
                    browse_name="DelmiaReceiver",
                    parent_gobject_id=1,
                    attributes=[
                        galaxy_pb.GalaxyAttribute(
                            attribute_name="DownloadPath",
                            full_tag_reference="DelmiaReceiver_001.DownloadPath",
                            mx_data_type=8,
                            data_type_name="String",
                        ),
                    ],
                ),
            ],
        ),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    objects = await client.discover_hierarchy()

    assert isinstance(objects, list)
    assert len(objects) == 2
    assert len(stub.discover_hierarchy.requests) == 2
    assert stub.discover_hierarchy.requests[0].page_size == 5000
    assert stub.discover_hierarchy.requests[0].page_token == ""
    assert stub.discover_hierarchy.requests[1].page_token == "page-2"
    assert objects[0].tag_name == "TestMachine_001"
    assert objects[1].attributes[0].full_tag_reference == "DelmiaReceiver_001.DownloadPath"


@pytest.mark.asyncio
async def test_discover_hierarchy_rejects_repeated_page_token() -> None:
    stub = FakeGalaxyStub()
    stub.discover_hierarchy.replies = [
        galaxy_pb.DiscoverHierarchyReply(next_page_token="7:1"),
        galaxy_pb.DiscoverHierarchyReply(next_page_token="7:1"),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    with pytest.raises(Exception, match="repeated page token"):
        await client.discover_hierarchy()


@pytest.mark.asyncio
async def test_watch_deploy_events_yields_events_in_order() -> None:
    ts1 = Timestamp()
    ts1.FromDatetime(datetime(2025, 4, 1, 10, 0, 0, tzinfo=timezone.utc))
    ts2 = Timestamp()
    ts2.FromDatetime(datetime(2025, 4, 1, 11, 0, 0, tzinfo=timezone.utc))
    events = [
        galaxy_pb.DeployEvent(
            sequence=1,
            observed_at=ts1,
            time_of_last_deploy=ts1,
            time_of_last_deploy_present=True,
            object_count=10,
            attribute_count=42,
        ),
        galaxy_pb.DeployEvent(
            sequence=2,
            observed_at=ts2,
            time_of_last_deploy=ts2,
            time_of_last_deploy_present=True,
            object_count=11,
            attribute_count=45,
        ),
    ]
    stub = FakeGalaxyStub()
    stub.watch_deploy_events.replies = list(events)
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    received: list[DeployEvent] = []
    async for event in client.watch_deploy_events():
        received.append(event)

    assert len(received) == 2
    assert received[0].sequence == 1
    assert received[1].sequence == 2
    assert received[0].object_count == 10
    assert received[1].attribute_count == 45
    assert stub.watch_deploy_events.metadata == (("authorization", "Bearer mxgw_test_secret"),)
    assert isinstance(stub.watch_deploy_events.requests[0], galaxy_pb.WatchDeployEventsRequest)
    # No last_seen_deploy_time was passed, so the request should leave it unset.
    assert not stub.watch_deploy_events.requests[0].HasField("last_seen_deploy_time")


@pytest.mark.asyncio
async def test_watch_deploy_events_propagates_last_seen_deploy_time() -> None:
    last_seen = datetime(2025, 4, 1, 12, 0, 0, tzinfo=timezone.utc)
    stub = FakeGalaxyStub()
    stub.watch_deploy_events.replies = []
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    async for _ in client.watch_deploy_events(last_seen_deploy_time=last_seen):
        pass

    request = stub.watch_deploy_events.requests[0]
    assert isinstance(request, WatchDeployEventsRequest)
    assert request.HasField("last_seen_deploy_time")
    assert request.last_seen_deploy_time.ToDatetime(tzinfo=timezone.utc) == last_seen


@pytest.mark.asyncio
async def test_watch_deploy_events_cancellation_closes_stream() -> None:
    ts = Timestamp()
    ts.FromDatetime(datetime(2025, 4, 1, 10, 0, 0, tzinfo=timezone.utc))
    stub = FakeGalaxyStub()
    # Use a "blocking" stream that never yields more after the first event.
    stub.watch_deploy_events = FakeStream(
        [galaxy_pb.DeployEvent(sequence=1, observed_at=ts)],
        block_after_replies=True,
    )
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    iterator = client.watch_deploy_events()

    first = await iterator.__anext__()
    assert first.sequence == 1

    # Break the iterator by aclose() — this should drive the cancel path.
    await iterator.aclose()

    assert stub.watch_deploy_events.cancel_called is True


@pytest.mark.asyncio
async def test_close_marks_channel_closed_when_no_real_channel() -> None:
    stub = FakeGalaxyStub()
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    await client.close()
    # Idempotent: a second close should not raise.
    await client.close()


def _obj(gid: int, tag: str, is_area: bool = False) -> galaxy_pb.GalaxyObject:
    return galaxy_pb.GalaxyObject(
        gobject_id=gid, tag_name=tag, browse_name=tag, is_area=is_area,
    )


def _build_browse_reply(
    children: list[galaxy_pb.GalaxyObject],
    child_has_children: list[bool],
    cache_sequence: int,
    next_page_token: str = "",
) -> galaxy_pb.BrowseChildrenReply:
    reply = galaxy_pb.BrowseChildrenReply(
        total_child_count=len(children),
        cache_sequence=cache_sequence,
        next_page_token=next_page_token,
    )
    reply.children.extend(children)
    reply.child_has_children.extend(child_has_children)
    return reply


def _fake_aio_rpc_error(code: grpc.StatusCode, details: str) -> grpc.aio.AioRpcError:
    return grpc.aio.AioRpcError(
        code=code,
        initial_metadata=grpc.aio.Metadata(),
        trailing_metadata=grpc.aio.Metadata(),
        details=details,
    )


@pytest.mark.asyncio
async def test_browse_no_parent_returns_roots() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply(
            children=[_obj(1, "Area_A", is_area=True), _obj(2, "Area_B", is_area=True)],
            child_has_children=[True, False],
            cache_sequence=7,
        ),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    roots = await client.browse()

    assert len(roots) == 2
    assert all(isinstance(node, LazyBrowseNode) for node in roots)
    assert roots[0].object.tag_name == "Area_A"
    assert roots[0].has_children_hint is True
    assert roots[1].has_children_hint is False
    assert roots[0].is_expanded is False
    request = stub.browse_children.requests[0]
    assert request.WhichOneof("parent") is None
    assert request.page_size == 500
    assert request.page_token == ""


@pytest.mark.asyncio
async def test_browse_expand_populates_children_and_marks_expanded() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply(
            children=[_obj(1, "Area_A", is_area=True)],
            child_has_children=[True],
            cache_sequence=1,
        ),
        _build_browse_reply(
            children=[_obj(11, "Child_A"), _obj(12, "Child_B")],
            child_has_children=[False, False],
            cache_sequence=1,
        ),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    roots = await client.browse()
    await roots[0].expand()

    assert roots[0].is_expanded is True
    assert [n.object.tag_name for n in roots[0].children] == ["Child_A", "Child_B"]
    assert len(stub.browse_children.requests) == 2
    expand_request = stub.browse_children.requests[1]
    assert expand_request.WhichOneof("parent") == "parent_gobject_id"
    assert expand_request.parent_gobject_id == 1


@pytest.mark.asyncio
async def test_browse_expand_idempotent_no_second_rpc() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply(
            children=[_obj(1, "Area_A", is_area=True)],
            child_has_children=[True],
            cache_sequence=1,
        ),
        _build_browse_reply(
            children=[_obj(11, "Child_A")],
            child_has_children=[False],
            cache_sequence=1,
        ),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    roots = await client.browse()
    await roots[0].expand()
    await roots[0].expand()

    assert len(stub.browse_children.requests) == 2
    assert len(roots[0].children) == 1


@pytest.mark.asyncio
async def test_browse_expand_concurrent_callers_only_fire_one_rpc() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply([_obj(1, "Plant", is_area=True)], [True], 7),
        _build_browse_reply([_obj(2, "Mixer_001")], [False], 7),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    roots = await client.browse()
    # Ten concurrent expand calls on the same node should issue exactly one RPC.
    await asyncio.gather(*(roots[0].expand() for _ in range(10)))

    assert roots[0].is_expanded
    assert len(roots[0].children) == 1
    # 1 roots fetch + exactly 1 expand fetch = 2 total
    assert len(stub.browse_children.requests) == 2


@pytest.mark.asyncio
async def test_browse_expand_unknown_parent_raises_mxgateway_error() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply(
            children=[_obj(99, "Stale_Parent", is_area=True)],
            child_has_children=[True],
            cache_sequence=1,
        ),
    ]
    stub.browse_children.exceptions = [
        None,
        _fake_aio_rpc_error(grpc.StatusCode.NOT_FOUND, "parent not found"),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    roots = await client.browse()
    with pytest.raises(MxGatewayError):
        await roots[0].expand()


@pytest.mark.asyncio
async def test_browse_expand_multi_page_gathers_all_pages() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply(
            children=[_obj(7, "Area_Big", is_area=True)],
            child_has_children=[True],
            cache_sequence=2,
        ),
        _build_browse_reply(
            children=[_obj(71, "Child_1"), _obj(72, "Child_2")],
            child_has_children=[False, False],
            cache_sequence=2,
            next_page_token="7:abc:2",
        ),
        _build_browse_reply(
            children=[_obj(73, "Child_3")],
            child_has_children=[False],
            cache_sequence=2,
        ),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )

    roots = await client.browse()
    await roots[0].expand()

    assert [n.object.tag_name for n in roots[0].children] == ["Child_1", "Child_2", "Child_3"]
    assert len(stub.browse_children.requests) == 3
    assert stub.browse_children.requests[2].page_token == "7:abc:2"
    assert stub.browse_children.requests[2].parent_gobject_id == 7


@pytest.mark.asyncio
async def test_browse_with_filter_forwards_to_request() -> None:
    stub = FakeGalaxyStub()
    stub.browse_children.replies = [
        _build_browse_reply(
            children=[_obj(1, "Area_A", is_area=True)],
            child_has_children=[False],
            cache_sequence=3,
        ),
    ]
    client = await GalaxyRepositoryClient.connect(
        ClientOptions(endpoint="fake", plaintext=True),
        stub=stub,
    )
    options = BrowseChildrenOptions(
        category_ids=(4, 5),
        template_chain_contains=("$DelmiaReceiver",),
        tag_name_glob="Area_*",
        include_attributes=True,
        alarm_bearing_only=True,
        historized_only=True,
    )

    await client.browse(options)

    request = stub.browse_children.requests[0]
    assert list(request.category_ids) == [4, 5]
    assert list(request.template_chain_contains) == ["$DelmiaReceiver"]
    assert request.tag_name_glob == "Area_*"
    assert request.HasField("include_attributes")
    assert request.include_attributes is True
    assert request.alarm_bearing_only is True
    assert request.historized_only is True


@pytest.mark.asyncio
async def test_browse_children_raw_returns_reply_unwrapped() -> None:
    """browse_children_raw forwards the request to the stub and returns the raw reply."""
    stub = FakeGalaxyStub()
    expected = _build_browse_reply(
        children=[_obj(1, "Plant", is_area=True)],
        child_has_children=[True],
        cache_sequence=42,
    )
    stub.browse_children.replies = [expected]

    async with await GalaxyRepositoryClient.connect(
        endpoint="fake",
        plaintext=True,
        stub=stub,
    ) as client:
        request = galaxy_pb.BrowseChildrenRequest(
            page_size=10,
            tag_name_glob="Plant*",
        )
        reply = await client.browse_children_raw(request)

    assert reply.cache_sequence == 42
    assert len(reply.children) == 1
    assert reply.children[0].tag_name == "Plant"
    assert len(stub.browse_children.requests) == 1
    assert stub.browse_children.requests[0].tag_name_glob == "Plant*"


class FakeGalaxyStub:
    def __init__(self) -> None:
        self.test_connection = FakeUnary([galaxy_pb.TestConnectionReply(ok=False)])
        self.get_last_deploy_time = FakeUnary([galaxy_pb.GetLastDeployTimeReply(present=False)])
        self.discover_hierarchy = FakeUnary([galaxy_pb.DiscoverHierarchyReply()])
        self.browse_children = FakeUnary([galaxy_pb.BrowseChildrenReply()])
        self.watch_deploy_events = FakeStream([])
        self.TestConnection = self.test_connection
        self.GetLastDeployTime = self.get_last_deploy_time
        self.DiscoverHierarchy = self.discover_hierarchy
        self.BrowseChildren = self.browse_children

    @property
    def WatchDeployEvents(self) -> "FakeStream":  # noqa: N802 — gRPC naming
        return self.watch_deploy_events


class FakeUnary:
    def __init__(self, replies: list[Any]) -> None:
        self.replies = replies
        self.requests: list[Any] = []
        # None entries mean "no exception on this call"; aligns with the replies queue index-by-index.
        self.exceptions: list[BaseException | None] = []
        self.metadata: tuple[tuple[str, str], ...] | None = None

    async def __call__(
        self,
        request: Any,
        *,
        metadata: tuple[tuple[str, str], ...],
        timeout: float | None = None,
    ) -> Any:
        self.requests.append(request)
        self.metadata = metadata
        if self.exceptions:
            exc = self.exceptions.pop(0)
            if exc is not None:
                raise exc
        return self.replies.pop(0)


class FakeStream:
    """Sync-callable fake matching the gRPC unary-stream surface.

    Calling the stub returns ``self`` (an async iterator). After exhausting the
    seeded ``replies``, iteration either ends (default) or blocks indefinitely
    (``block_after_replies=True``) so cancellation paths can be exercised.
    """

    def __init__(
        self,
        replies: list[Any],
        *,
        block_after_replies: bool = False,
    ) -> None:
        self.replies = list(replies)
        self.requests: list[Any] = []
        self.metadata: tuple[tuple[str, str], ...] | None = None
        self.cancel_called = False
        self._block_after_replies = block_after_replies

    def __call__(
        self,
        request: Any,
        *,
        metadata: tuple[tuple[str, str], ...],
        timeout: float | None = None,
    ) -> "FakeStream":
        self.requests.append(request)
        self.metadata = metadata
        return self

    def __aiter__(self) -> "FakeStream":
        return self

    async def __anext__(self) -> Any:
        if self.replies:
            return self.replies.pop(0)
        if self._block_after_replies:
            # Sleep forever until the consumer cancels us.
            await asyncio.Event().wait()
        raise StopAsyncIteration

    def cancel(self) -> None:
        self.cancel_called = True
