import unittest
from sia.llm_engine import LlmEngine, InferenceResult

echo_system_prompt = """
    Your answer always consists of 4 parts:
        - The original request
        - <test_tag>
        - The original request
        - </test_tag>
        You never provide an answer.
        Only the entered text, the xml open tag, the same text again and the xml close tag.
        Don't add whitespace or newlines.
        Don't modify the text or change casing.
        Be exact, this is for testing purposes.
    """

echo_action_schema = """
    <?xml version="1.0" encoding="UTF-8"?>
    <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
    <xs:element name="test_tag">
        <xs:complexType>
        <xs:simpleContent>
            <xs:extension base="xs:string">
            <xs:attribute name="id" type="xs:string" use="required"/>
            </xs:extension>
        </xs:simpleContent>
        </xs:complexType>
    </xs:element>
    </xs:schema>
    """

class TestLlmEngine(unittest.TestCase):
    def setUp(self):
        self.model_path = "/root/model"

    def test_initialization(self):
        llm_engine = LlmEngine(self.model_path)
        self.assertIsInstance(llm_engine, LlmEngine)

    def test_infer(self):
        main_context = "This is a test"
        llm_engine = LlmEngine(self.model_path)
        result = llm_engine.infer(echo_system_prompt, main_context, echo_action_schema)
        self.assertIsInstance(result, InferenceResult)
        self.assertIsInstance(result.reasoning, str)
        self.assertIsInstance(result.actions, str)
        self.assertEqual(result.reasoning, main_context)
        self.assertEqual(result.actions, f"<test_tag>{main_context}</test_tag>")

'''
    def test_set_model_path(self):
        new_model_path = "/path/to/new/model"
        self.llm_engine.set_model_path(new_model_path)
        # Add assertions to check if the model path was updated correctly

    def test_finetune(self):
        dataset_paths = ["/path/to/dataset1", "/path/to/dataset2"]
        output_dir = "/path/to/output"

        self.llm_engine.finetune(dataset_paths, output_dir)
        # Add assertions to check if the fine-tuning process completed successfully
        # For example, check if new model weights were saved in the output directory
'''

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