from __future__ import annotations

import unittest
from unittest.mock import patch

from Mama.models import ConversationState, ExecutionContext, LlmConfig, RetrievalConfig, StorageConfig
from Mama.runtime import DefaultLlmProvider


class RuntimeOllamaFormatTest(unittest.TestCase):
    def _ctx(self, parameters):
        return ExecutionContext(
            tenant_id="t",
            kb_id="k",
            user_id="u",
            session_id="s",
            llm=LlmConfig(provider="ollama", model="qwen2.5:3b", parameters=parameters),
            retrieval=RetrievalConfig(search_type="similarity", k=2),
            storage=StorageConfig(backend="faiss_local", params={}),
            conversation=ConversationState(messages=[], max_messages=5),
        )

    def test_sanitize_ollama_params_accepts_schema_object(self):
        schema = {"type": "object", "properties": {"title": {"type": "string"}}}
        out = DefaultLlmProvider._sanitize_ollama_params({"format": schema, "temperature": 0.1})
        self.assertEqual(out["format"], schema)
        self.assertEqual(out["temperature"], 0.1)

    def test_build_llm_passes_schema_format_to_chat_ollama(self):
        schema = {"type": "object", "required": ["x"], "properties": {"x": {"type": "string"}}}
        captured = {}

        def _fake_chat_ollama(**kwargs):
            captured.update(kwargs)
            return object()

        with patch("Mama.runtime.ChatOllama", new=_fake_chat_ollama):
            provider = DefaultLlmProvider()
            provider.build_llm(self._ctx({"format": schema, "temperature": 0.2}))

        self.assertEqual(captured.get("format"), schema)
        self.assertEqual(captured.get("temperature"), 0.2)
        self.assertEqual(captured.get("model"), "qwen2.5:3b")


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