import io
import json
import unittest
import urllib.error
from dataclasses import asdict
from typing import Any, Dict
from unittest.mock import patch

# Adjust these imports to your real module paths
from src.aiel_sdk.services.memory import MemoryService
from src.aiel_sdk.models.memory import (
    ThreadAppendIn,
    ThreadScope,
    ChatMessageIn,
    ThreadGetOut,
    ThreadStatePatchIn,
    CheckpointPutIn,
    CheckpointOut,
    RecallForgetIn,
    ThreadStateOut
)
from src.aiel_sdk.memory.in_memory import InMemorySaver
from src.aiel_sdk.memory.remote import RemoteSaver
from src.aiel_sdk.memory.types import ThreadMessage, ThreadScope as SaverThreadScope, RecallQuery


class _FakeResponse:
    def __init__(self, body: bytes):
        self._body = body

    def read(self) -> bytes:
        return self._body

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class _CaptureRequestUrlopen:
    """
    Captures urllib.request.Request objects passed to urlopen.
    """
    def __init__(self, response_body: bytes):
        self.response_body = response_body
        self.last_request = None

    def __call__(self, req, timeout=None):
        self.last_request = req
        return _FakeResponse(self.response_body)


class _FakeAielClient:
    """
    Stub client used by MemoryService tests.
    It mimics AielClient._request(...) signature.
    """
    def __init__(self, result: Any = None):
        self.calls = []
        self.result = result

    def _request(self, service: str, method: str, path: str, *, params=None, json_body=None, request_id=None):
        self.calls.append(
            {
                "service": service,
                "method": method,
                "path": path,
                "params": params,
                "json_body": json_body,
                "request_id": request_id,
            }
        )
        return self.result


class MemoryModelsTests(unittest.TestCase):
    def test_chat_message_role_is_restricted(self):
        # role must be one of human/ai/system/tool
        with self.assertRaises(Exception):
            ChatMessageIn(role="user", content="hello")  # invalid

        ok = ChatMessageIn(role="human", content="hello")
        self.assertEqual(ok.role, "human")

    def test_chat_message_content_limits(self):
        with self.assertRaises(Exception):
            ChatMessageIn(role="human", content="")  # min_length=1

        big = "x" * 20001
        with self.assertRaises(Exception):
            ChatMessageIn(role="human", content=big)

    def test_thread_append_constraints(self):
        with self.assertRaises(Exception):
            ThreadAppendIn(scope=ThreadScope(), messages=[])  # min_items=1

        msgs = [ChatMessageIn(role="human", content="hi")]
        payload = ThreadAppendIn(scope=ThreadScope(user_id="u1"), messages=msgs, ttl_seconds=60)
        self.assertEqual(payload.scope.user_id, "u1")
        self.assertEqual(len(payload.messages), 1)


class MemoryServiceTests(unittest.TestCase):
    def test_thread_key_is_url_quoted(self):
        """
        Thread key might contain spaces or slashes; it must be URL encoded.
        """
        fake = _FakeAielClient(result={"thread_id": "th_abc", "messages": [], "updated_at": "2026-03-01T00:00:00Z"})
        svc = MemoryService(fake)

        thread_key = "support/ticket 123"  # contains slash + space
        payload = ThreadAppendIn(
            scope=ThreadScope(),
            messages=[ChatMessageIn(role="human", content="hello")],
            ttl_seconds=60,
        )

        out = svc.thread_append_messages("w1", "p1", thread_key=thread_key, payload=payload)
        self.assertEqual(out.thread_id, "th_abc")

        call = fake.calls[-1]
        self.assertEqual(call["service"], "memory")
        self.assertEqual(call["method"], "POST")
        # slash and space should be encoded in the path segment
        self.assertIn("/memory/thread/support%2Fticket%20123/messages:append", call["path"])

    def test_thread_get_messages_uses_limit_param_only(self):
        fake = _FakeAielClient(result={"thread_id": "th_x", "messages": [], "updated_at": "2026-03-01T00:00:00Z"})
        svc = MemoryService(fake)

        out = svc.thread_get_messages("w1", "p1", thread_key="t1", limit=77)
        self.assertEqual(out.thread_id, "th_x")

        call = fake.calls[-1]
        self.assertEqual(call["method"], "GET")
        self.assertEqual(call["params"], {"limit": 77})

    def test_state_patch_serializes_pydantic_payload(self):
        fake = _FakeAielClient(result={"thread_id": "th_x", "state": {"a": 1}, "updated_at": "2026-03-01T00:00:00Z"})
        svc = MemoryService(fake)

        payload = ThreadStatePatchIn(scope=ThreadScope(user_id="u1"), patch={"a": 1}, ttl_seconds=60)
        out = svc.thread_patch_state("w1", "p1", thread_key="t1", payload=payload)

        self.assertEqual(out.state["a"], 1)
        call = fake.calls[-1]
        self.assertEqual(call["json_body"]["scope"]["user_id"], "u1")
        self.assertEqual(call["json_body"]["patch"]["a"], 1)

    def test_checkpoint_returns_typed_output(self):
        fake = _FakeAielClient(
            result={
                "thread_id": "th_x",
                "checkpoint_id": "cp_1",
                "state": {"step": "triage"},
                "metadata": {"node": "triage"},
                "created_at": "2026-03-01T00:00:00Z",
            }
        )
        svc = MemoryService(fake)

        payload = CheckpointPutIn(scope=ThreadScope(agent_id="a1"), state={"step": "triage"}, metadata={"node": "triage"}, ttl_seconds=60)
        out = svc.checkpoint_put("w1", "p1", thread_key="t1", payload=payload)

        # If you kept checkpoint_put returning CheckpointOut in your service, this passes:
        # If you kept dict return, update accordingly.
        self.assertTrue(hasattr(out, "checkpoint_id"))
        self.assertEqual(out.checkpoint_id, "cp_1")

    def test_recall_forget_accepts_typed_payload(self):
        fake = _FakeAielClient(result={"ok": True})
        svc = MemoryService(fake)

        payload = RecallForgetIn(namespace=["proj", "user"], selector={"key": "k1"})
        out = svc.recall_forget("w1", "p1", payload=payload)
        self.assertEqual(out.get("ok"), True)

        call = fake.calls[-1]
        self.assertEqual(call["json_body"]["namespace"], ["proj", "user"])
        self.assertEqual(call["json_body"]["selector"]["key"], "k1")


class InMemorySaverTests(unittest.TestCase):
    def test_inmemory_thread_write_and_read(self):
        saver = InMemorySaver()

        saver.thread_write(
            thread_key="t1",
            messages=[ThreadMessage(role="human", content="hello")],
            scope=SaverThreadScope(user_id="u1"),
            ttl_seconds=60,
        )

        t = saver.thread_read(thread_key="t1", limit=10)
        self.assertEqual(t.thread_key, "t1")
        self.assertEqual(len(t.messages), 1)
        self.assertEqual(t.messages[0].content, "hello")

    def test_inmemory_state_patch(self):
        saver = InMemorySaver()
        saver.state_patch(thread_key="t1", patch={"x": 1})
        st = saver.state_read(thread_key="t1")
        self.assertEqual(st.state["x"], 1)

    def test_inmemory_checkpoint_latest(self):
        saver = InMemorySaver()
        self.assertIsNone(saver.checkpoint_read_latest(thread_key="t1"))
        saver.checkpoint_write(thread_key="t1", state={"s": 1}, metadata={"m": 1})
        latest = saver.checkpoint_read_latest(thread_key="t1")
        self.assertIsNotNone(latest)
        self.assertEqual(latest.state["s"], 1)

    def test_inmemory_recall(self):
        saver = InMemorySaver()
        saver.recall_put(namespace=["a"], key="k1", value={"v": 1})
        item = saver.recall_get(namespace=["a"], key="k1")
        self.assertIsNotNone(item)
        self.assertEqual(item.value["v"], 1)

        items = saver.recall_search(query=RecallQuery(namespace=["a"], limit=10))
        self.assertEqual(len(items), 1)

        saver.recall_forget_key(namespace=["a"], key="k1")
        self.assertIsNone(saver.recall_get(namespace=["a"], key="k1"))


class RemoteSaverTests(unittest.TestCase):
    def test_remote_saver_calls_memory_service_correctly(self):
        """
        This verifies RemoteSaver builds correct payloads and calls MemoryService methods.
        """
        class _FakeMemory:
            def __init__(self):
                self.calls = []

            def thread_append_messages(self, ws, proj, *, thread_key, payload):
                self.calls.append(("append", ws, proj, thread_key, payload))
                return ThreadGetOut(thread_id="th_x", messages=[], updated_at="2026-03-01T00:00:00Z")

            def thread_get_messages(self, ws, proj, *, thread_key, limit=50):
                self.calls.append(("get", ws, proj, thread_key, limit))
                return ThreadGetOut(thread_id="th_x", messages=[], updated_at="2026-03-01T00:00:00Z")

            def thread_get_state(self, ws, proj, *, thread_key):
                self.calls.append(("state_get", ws, proj, thread_key))
                return ThreadStateOut(thread_id="th_x", state={"a": 1}, updated_at="2026-03-01T00:00:00Z")

            def thread_patch_state(self, ws, proj, *, thread_key, payload):
                self.calls.append(("state_patch", ws, proj, thread_key, payload))
                return ThreadStateOut(thread_id="th_x", state={"a": 2}, updated_at="2026-03-01T00:00:00Z")

            def checkpoint_get_latest(self, ws, proj, *, thread_key):
                self.calls.append(("cp_latest", ws, proj, thread_key))
                return {
                    "thread_id": "th_x",
                    "checkpoint_id": "cp_1",
                    "state": {"s": 1},
                    "metadata": {"m": 1},
                    "created_at": "2026-03-01T00:00:00Z",
                }

            def checkpoint_put(self, ws, proj, *, thread_key, payload):
                self.calls.append(("cp_put", ws, proj, thread_key, payload))
                return {
                    "thread_id": "th_x",
                    "checkpoint_id": "cp_2",
                    "state": payload.get("state") or {},
                    "metadata": payload.get("metadata") or {},
                    "created_at": "2026-03-01T00:00:00Z",
                }

            def recall_put(self, ws, proj, *, payload):
                self.calls.append(("recall_put", ws, proj, payload))
                return {"ok": True}

            def recall_get(self, ws, proj, *, payload):
                self.calls.append(("recall_get", ws, proj, payload))
                return {"namespace": payload["namespace"], "key": payload["key"], "value": {"v": 1}, "tags": {}}

            def recall_search(self, ws, proj, *, payload):
                self.calls.append(("recall_search", ws, proj, payload))
                return {"items": [{"namespace": payload["namespace"], "key": "k1", "value": {"v": 1}, "tags": {}}]}

            def recall_forget(self, ws, proj, *, payload):
                self.calls.append(("recall_forget", ws, proj, payload))
                return {"ok": True}

        class _FakeClient:
            def __init__(self):
                self.memory = _FakeMemory()

        client = _FakeClient()
        saver = RemoteSaver(client, workspace_id="w1", project_id="p1")

        # write message
        saver.thread_write(thread_key="t1", messages=[ThreadMessage(role="human", content="hello")])
        self.assertEqual(client.memory.calls[0][0], "append")
        self.assertEqual(client.memory.calls[0][3], "t1")
        sent_payload = client.memory.calls[0][4]
        self.assertIn("scope", sent_payload)
        self.assertEqual(sent_payload["messages"][0]["role"], "human")

        # patch state
        st = saver.state_patch(thread_key="t1", patch={"a": 2})
        self.assertEqual(st.state["a"], 2)
        self.assertEqual(client.memory.calls[1][0], "state_patch")

        # checkpoint
        cp = saver.checkpoint_write(thread_key="t1", state={"s": 2}, metadata={"m": 2})
        self.assertEqual(cp.checkpoint_id, "cp_2")
        self.assertEqual(client.memory.calls[2][0], "cp_put")

        # recall
        saver.recall_put(namespace=["n"], key="k", value={"v": 1})
        item = saver.recall_get(namespace=["n"], key="k")
        self.assertEqual(item.value["v"], 1)
        items = saver.recall_search(query=RecallQuery(namespace=["n"], limit=10))
        self.assertEqual(len(items), 1)
        saver.recall_forget_key(namespace=["n"], key="k")
        self.assertEqual(client.memory.calls[-1][0], "recall_forget")


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