from __future__ import annotations

import importlib
import sys
import types
import unittest
from pathlib import Path


def install_langchain_stub() -> None:
    """Install minimal langchain stubs so we can import mubit_langchain without langchain."""
    memory_module = types.ModuleType("langchain_core.memory")

    class BaseMemory:
        class Config:
            arbitrary_types_allowed = True

    memory_module.BaseMemory = BaseMemory

    messages_module = types.ModuleType("langchain_core.messages")

    class SystemMessage:
        def __init__(self, content=""):
            self.content = content

    messages_module.SystemMessage = SystemMessage

    langchain_core = types.ModuleType("langchain_core")
    langchain_core.memory = memory_module
    langchain_core.messages = messages_module

    sys.modules["langchain_core"] = langchain_core
    sys.modules["langchain_core.memory"] = memory_module
    sys.modules["langchain_core.messages"] = messages_module


# Install stubs immediately so that mubit_langchain can be imported
# during unittest discovery (it imports langchain_core at module level).
install_langchain_stub()


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

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

    def recall(self, **payload):
        self.calls.append(("recall", payload))
        return {
            "evidence": [
                {
                    "id": "mem-1",
                    "content": "Relevant context",
                    "metadata_json": '{"source":"langchain"}',
                    "score": 0.9,
                }
            ]
        }

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

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


def _import_module():
    package_root = Path(__file__).resolve().parents[2]
    if str(package_root) not in sys.path:
        sys.path.insert(0, str(package_root))
    # Clear cached imports to pick up stubs
    for mod_name in list(sys.modules):
        if mod_name.startswith("mubit_langchain"):
            del sys.modules[mod_name]
    return importlib.import_module("mubit_langchain")


class TestMubitMemory(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.module = _import_module()

    def _make_memory(self, client):
        mem = self.module.MubitMemory.__new__(self.module.MubitMemory)
        mem.endpoint = "http://127.0.0.1:3000"
        mem.api_key = ""
        mem.session_id = "chain-1"
        mem.user_id = ""
        mem.agent_id = "langchain"
        mem.memory_key = "history"
        mem.input_key = None
        mem.output_key = None
        mem.return_messages = False
        mem.max_token_budget = 2048
        mem.entry_types = None
        mem.context_sections = None
        mem._client = client
        return mem

    def test_memory_variables_returns_memory_key(self) -> None:
        client = MockClient()
        memory = self._make_memory(client)

        self.assertEqual(memory.memory_variables, ["history"])

    def test_load_memory_variables_fetches_context(self) -> None:
        client = MockClient()
        memory = self._make_memory(client)

        result = memory.load_memory_variables({"input": "What happened yesterday?"})

        self.assertEqual(result["history"], "assembled context from MuBit")
        self.assertEqual(len(client.calls), 1)
        self.assertEqual(client.calls[0][0], "get_context")
        self.assertEqual(client.calls[0][1]["session_id"], "chain-1")

    def test_load_memory_variables_empty_on_no_query(self) -> None:
        client = MockClient()
        memory = self._make_memory(client)

        result = memory.load_memory_variables({})

        self.assertEqual(result["history"], "")
        self.assertEqual(len(client.calls), 0)

    def test_save_context_ingests_interaction(self) -> None:
        client = MockClient()
        memory = self._make_memory(client)

        memory.save_context(
            {"input": "What is the weather?"},
            {"output": "It's sunny today."},
        )

        self.assertEqual(len(client.calls), 1)
        name, payload = client.calls[0]
        self.assertEqual(name, "remember")
        self.assertIn("What is the weather?", payload["text"])
        self.assertIn("It's sunny today.", payload["text"])
        self.assertEqual(payload["intent"], "trace")

    def test_clear_calls_delete_run(self) -> None:
        client = MockClient()
        memory = self._make_memory(client)

        memory.clear()

        self.assertEqual(len(client.calls), 1)
        self.assertEqual(client.calls[0][0], "delete_run")
        self.assertEqual(client.calls[0][1]["session_id"], "chain-1")

    def test_input_key_override(self) -> None:
        client = MockClient()
        memory = self._make_memory(client)
        memory.input_key = "question"

        memory.load_memory_variables({"question": "Hello", "other": "ignored"})

        self.assertEqual(client.calls[0][1]["query"], "Hello")


class TestMubitChatMemory(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.module = _import_module()

    def _make_chat_memory(self, client):
        mem = self.module.MubitChatMemory.__new__(self.module.MubitChatMemory)
        mem.endpoint = "http://127.0.0.1:3000"
        mem.api_key = ""
        mem.session_id = "chat-1"
        mem.user_id = ""
        mem.agent_id = "langchain"
        mem.memory_key = "history"
        mem.input_key = None
        mem.output_key = None
        mem.return_messages = True
        mem.max_token_budget = 2048
        mem.entry_types = None
        mem.context_sections = None
        mem._client = client
        return mem

    def test_save_context_stores_messages_individually(self) -> None:
        client = MockClient()
        memory = self._make_chat_memory(client)

        memory.save_context(
            {"input": "Hello"},
            {"output": "Hi there!"},
        )

        self.assertEqual(len(client.calls), 2)
        self.assertEqual(client.calls[0][1]["intent"], "context")
        self.assertEqual(client.calls[0][1]["text"], "Hello")
        self.assertEqual(client.calls[1][1]["intent"], "trace")
        self.assertEqual(client.calls[1][1]["text"], "Hi there!")

    def test_return_messages_returns_system_messages(self) -> None:
        client = MockClient()
        memory = self._make_chat_memory(client)

        result = memory.load_memory_variables({"input": "test query"})

        self.assertTrue(isinstance(result["history"], list))
        self.assertEqual(len(result["history"]), 1)
        self.assertEqual(result["history"][0].content, "assembled context from MuBit")


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