import unittest

from langchain.base_language import BaseLanguageModel
from langchain.chat_models.base import BaseChatModel
from langchain.prompts.chat import ChatPromptTemplate
from langchain.prompts.example_selector.base import BaseExampleSelector
from langchain.schema import AIMessage, BaseRetriever, HumanMessage

from .resolver import Resolver


class TestResolver(unittest.TestCase):
    def test_simple_load(self):
        r = Resolver("names.yaml")

    def test_names(self):
        r = Resolver("names.yaml")
        self.assertTrue(len(r.names()) > 0)

    def test_lookup_and_relookup(self):
        r = Resolver("names.yaml")

        cases = {
            "text": str,
            "float": float,
            "integer": int,
            "list": list,
        }
        for name, cls in cases.items():
            self.assertEqual(r.lookup(name), cls)
            self.assertEqual(r.relookup(cls), name)

    def test_abstract(self):
        r = Resolver("names.yaml")

        true_cases = [
            BaseChatModel,
            BaseLanguageModel,
            BaseExampleSelector,
            BaseRetriever,
        ]
        for c in true_cases:
            self.assertTrue(r.is_abstract(c))

        false_cases = [
            ChatPromptTemplate,
            AIMessage,
            HumanMessage,
        ]
        for c in false_cases:
            self.assertFalse(r.is_abstract(c))

    def test_candidates(self):
        r = Resolver("names.yaml")
        candidates = r.candidates("LLM")
        self.assertTrue("LLM" not in candidates)
        self.assertTrue("ChatLLM" not in candidates)

    def test_slots(self):
        r = Resolver("names.yaml")
        # print(r.slots('LLM'))


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