from __future__ import annotations

import importlib
import sys
import types
import unittest
from dataclasses import dataclass, field
from pathlib import Path


def install_google_adk_stub() -> None:
    google = types.ModuleType("google")
    adk = types.ModuleType("google.adk")
    memory = types.ModuleType("google.adk.memory")
    sessions = types.ModuleType("google.adk.sessions")
    events = types.ModuleType("google.adk.events")
    genai = types.ModuleType("google.genai")
    genai_types = types.ModuleType("google.genai.types")

    class BaseMemoryService:
        pass

    @dataclass
    class SearchMemoryResponse:
        memories: list

    @dataclass
    class Part:
        text: str = ""

    @dataclass
    class Content:
        parts: list[Part]
        role: str = "model"

    @dataclass
    class MemoryEntry:
        content: Content
        custom_metadata: dict | None = None
        id: str | None = None
        author: str | None = None

    @dataclass
    class Session:
        id: str
        user_id: str
        app_name: str
        events: list

    @dataclass
    class Event:
        author: str | None = None
        content: Content | None = None
        text: str | None = None

    memory.base_memory_service = types.ModuleType("google.adk.memory.base_memory_service")
    memory.base_memory_service.BaseMemoryService = BaseMemoryService
    memory.base_memory_service.SearchMemoryResponse = SearchMemoryResponse
    memory.memory_entry = types.ModuleType("google.adk.memory.memory_entry")
    memory.memory_entry.MemoryEntry = MemoryEntry
    sessions.session = types.ModuleType("google.adk.sessions.session")
    sessions.session.Session = Session
    events.event = types.ModuleType("google.adk.events.event")
    events.event.Event = Event

    genai.types = genai_types
    genai_types.Part = Part
    genai_types.Content = Content

    google.adk = adk
    google.genai = genai
    adk.memory = memory
    adk.sessions = sessions
    adk.events = events

    sys.modules["google"] = google
    sys.modules["google.adk"] = adk
    sys.modules["google.adk.memory"] = memory
    sys.modules["google.adk.memory.base_memory_service"] = memory.base_memory_service
    sys.modules["google.adk.memory.memory_entry"] = memory.memory_entry
    sys.modules["google.adk.sessions"] = sessions
    sys.modules["google.adk.sessions.session"] = sessions.session
    sys.modules["google.adk.events"] = events
    sys.modules["google.adk.events.event"] = events.event
    sys.modules["google.genai"] = genai
    sys.modules["google.genai.types"] = genai_types


class MockClient:
    def __init__(self) -> None:
        self.calls: list[tuple[str, dict]] = []

    def ingest_items(self, **payload):
        self.calls.append(("ingest_items", payload))
        return {"accepted": True, "job_id": "job-1"}

    def recall(self, **payload):
        self.calls.append(("recall", payload))
        return {
            "evidence": [
                {
                    "id": "memory-1",
                    "content": "Stored memory",
                    "metadata_json": '{"author":"planner","channel":"incident"}',
                    "score": 0.82,
                    "evidence_type": "lesson",
                }
            ]
        }

    def checkpoint(self, **payload):
        self.calls.append(("checkpoint", payload))
        return {"success": True}

    def record_outcome(self, **payload):
        self.calls.append(("record_outcome", payload))
        return {"success": True}

    def surface_strategies(self, **payload):
        self.calls.append(("surface_strategies", payload))
        return {"strategies": []}

    def get_context(self, **payload):
        self.calls.append(("get_context", payload))
        return {"context_block": "assembled"}

    def archive(self, **payload):
        self.calls.append(("archive", payload))
        return {"success": True, "reference_id": "archive-1"}

    def dereference(self, **payload):
        self.calls.append(("dereference", payload))
        return {"found": True, "evidence": {"id": "archive-1"}}

    def memory_health(self, **payload):
        self.calls.append(("memory_health", payload))
        return {"entry_counts": {"lesson": 1}}

    def diagnose(self, **payload):
        self.calls.append(("diagnose", payload))
        return {"failure_lessons": []}

    def register_agent(self, **payload):
        self.calls.append(("register_agent", payload))
        return {"success": True}

    def list_agents(self, **payload):
        self.calls.append(("list_agents", payload))
        return {"agents": []}

    def handoff(self, **payload):
        self.calls.append(("handoff", payload))
        return {"success": True, "handoff_id": "handoff-1"}

    def feedback(self, **payload):
        self.calls.append(("feedback", payload))
        return {"success": True, "feedback_id": "feedback-1"}


class TestMubitAdkMemoryService(unittest.IsolatedAsyncioTestCase):
    @classmethod
    def setUpClass(cls) -> None:
        install_google_adk_stub()
        package_root = Path(__file__).resolve().parents[2]
        if str(package_root) not in sys.path:
            sys.path.insert(0, str(package_root))
        cls.module = importlib.import_module("mubit_adk")
        cls.genai_types = sys.modules["google.genai.types"]
        cls.session_cls = sys.modules["google.adk.sessions.session"].Session
        cls.event_cls = sys.modules["google.adk.events.event"].Event

    async def test_add_events_to_memory_routes_through_sdk_client(self) -> None:
        client = MockClient()
        service = self.module.MubitMemoryService(mubit_client=client)
        event = self.event_cls(
            author="tool",
            content=self.genai_types.Content(parts=[self.genai_types.Part(text="built release manifest")]),
        )

        await service.add_events_to_memory(
            app_name="deploy-agent",
            user_id="user-1",
            events=[event],
            session_id="session-1",
            custom_metadata={"phase": "release"},
        )

        self.assertEqual(
            client.calls[0],
            (
                "ingest_items",
                {
                    "session_id": "session-1",
                    "agent_id": "deploy-agent",
                    "idempotency_key": "",
                    "items": [
                        {
                            "item_id": "adk-evt-0",
                            "content_type": "text/plain",
                            "text": "built release manifest",
                            "metadata": {
                                "author": "tool",
                                "event_index": 0,
                                "phase": "release",
                            },
                            "user_id": "user-1",
                            "intent": "tool_output",
                        }
                    ],
                },
            ),
        )

    async def test_search_memory_normalizes_metadata_json(self) -> None:
        client = MockClient()
        service = self.module.MubitMemoryService(mubit_client=client)

        response = await service.search_memory(
            app_name="deploy-agent",
            user_id="user-1",
            query="release memory",
        )

        self.assertEqual(len(response.memories), 1)
        memory = response.memories[0]
        self.assertEqual(memory.author, "planner")
        self.assertEqual(memory.custom_metadata["channel"], "incident")
        self.assertEqual(memory.custom_metadata["score"], 0.82)

    async def test_current_helper_surface_is_available(self) -> None:
        client = MockClient()
        service = self.module.MubitMemoryService(mubit_client=client)

        await service.get_context(
            app_name="deploy-agent",
            user_id="user-1",
            query="release status",
            session_id="session-1",
        )
        await service.archive(
            app_name="deploy-agent",
            user_id="user-1",
            session_id="session-1",
            content="pytest failure summary",
            artifact_kind="failing_test_summary",
            source_tool="pytest",
        )
        await service.dereference(
            user_id="user-1",
            session_id="session-1",
            reference_id="archive-1",
        )
        await service.memory_health(
            user_id="user-1",
            session_id="session-1",
            limit=50,
        )
        await service.diagnose(
            user_id="user-1",
            error_text="kubectl apply failed",
            session_id="session-1",
        )
        await service.register_agent(
            user_id="user-1",
            session_id="session-1",
            agent_id="planner",
            role="planner",
        )
        await service.list_agents(user_id="user-1", session_id="session-1")
        await service.handoff(
            user_id="user-1",
            session_id="session-1",
            from_agent_id="planner",
            to_agent_id="reviewer",
            content="Review the release notes.",
        )
        await service.feedback(
            user_id="user-1",
            session_id="session-1",
            handoff_id="handoff-1",
            verdict="approve",
        )

        self.assertEqual(
            [name for name, _payload in client.calls],
            [
                "get_context",
                "archive",
                "dereference",
                "memory_health",
                "diagnose",
                "register_agent",
                "list_agents",
                "handoff",
                "feedback",
            ],
        )


if __name__ == "__main__":
    unittest.main()
