import os
import shutil
import tempfile
import unittest

import transformers
from packaging import version

from swift.llm import ExportArguments, export_main

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'


class TestTemplate(unittest.TestCase):

    def setUp(self):
        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
        self.tmp_dir = tempfile.TemporaryDirectory().name

    def tearDown(self):
        if os.path.exists(self.tmp_dir):
            shutil.rmtree(self.tmp_dir)
        super().tearDown()

    @unittest.skip('swift2.0')
    def test_llama3(self):
        args = ExportArguments(model_type='llama3-8b-instruct', to_ollama=True, ollama_output_dir=self.tmp_dir)
        export_main(args)

        template = ('TEMPLATE """{{ if .System }}<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n'
                    '{{ .System }}<|eot_id|>{{ else }}<|begin_of_text|>{{ end }}{{ if .Prompt }}<|start_header_id|>user'
                    '<|end_header_id|>\n\n{{ .Prompt }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n'
                    '{{ end }}{{ .Response }}<|eot_id|>"""')

        stop = 'PARAMETER stop "<|eot_id|>"'

        with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
            content = f.read()
            self.assertTrue(template in content)
            self.assertTrue(stop in content)

    @unittest.skip('swift2.0')
    def test_glm4(self):
        if version.parse(transformers.__version__) >= version.parse('4.45'):
            return

        args = ExportArguments(model_type='glm4-9b-chat', to_ollama=True, ollama_output_dir=self.tmp_dir)
        export_main(args)

        template = ('TEMPLATE """{{ if .System }}[gMASK] <sop><|system|>\n{{ .System }}{{ else }}'
                    '[gMASK] <sop>{{ end }}{{ if .Prompt }}<|user|>\n{{ .Prompt }}<|assistant|>\n'
                    '{{ end }}{{ .Response }}<|user|>"""')

        stop = 'PARAMETER stop "<|user|>"'

        with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
            content = f.read()
            self.assertTrue(template in content)
            self.assertTrue(stop in content)

    @unittest.skip('swift2.0')
    def test_qwen2(self):
        args = ExportArguments(model_type='qwen2-7b-instruct', to_ollama=True, ollama_output_dir=self.tmp_dir)
        export_main(args)

        template = ('TEMPLATE """{{ if .System }}<|im_start|>system\n{{ .System }}<|im_end|>\n{{ else }}{{ end }}'
                    '{{ if .Prompt }}<|im_start|>user\n{{ .Prompt }}<|im_end|>\n<|im_start|>assistant\n'
                    '{{ end }}{{ .Response }}<|im_end|>"""')

        stop = 'PARAMETER stop "<|im_end|>"'

        with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
            content = f.read()
            self.assertTrue(template in content)
            self.assertTrue(stop in content)


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