import json
import os


class CodeGenerator:
    def __init__(self):
        self.import_str = 'from manimlib.imports import *'
        self.class_declare_template = 'class {0}(Scene):'
        self.construct_str = 'def construct(self):'
        self.indent = '    '
        self.code_list = []
        self.wait_str = 'self.wait()'

    def assemble_code(self, codes):
        self.code_list.append(self.import_str)
        class_declare = self.class_declare_template.format(codes['scene_name'])
        self.code_list.append(class_declare)
        self.code_list.append(self.indent + self.construct_str)
        indent_cur = 2
        for item in codes['codes']:
            param_list = []
            for param in item['params']:
                if param == 'varargs':
                    vararg_list = []
                    for sub_param in item['params'][param]:
                        if sub_param['type'] == 'string':
                            vararg_list.append('\"' + sub_param['content'] + '\"')
                        elif sub_param['type'] == 'instance':
                            vararg_list.append(sub_param['content'])
                    param_list.extend(vararg_list)
                else:
                    if item['params'][param]['type'] == 'string':
                        param_display = '\'' + item['params'][param]['content'] + '\''
                    elif item['params'][param]['type'] == 'constant':
                        param_display = item['params'][param]['content']
                    else:
                        param_display = str(item['params'][param]['content'])
                    param_item = param + '=' + param_display
                    param_list.append(param_item)
            param_str = ', '.join(param_list)
            if item['type'] == 'create_instance':
                construct_str = item['class'] + '(' + param_str + ')'
                line = self.indent * indent_cur + item['name'] + ' = ' + construct_str
                self.code_list.append(line)
            elif item['type'] == 'call_function':
                function_str = item['function'] + '(' + param_str + ')'
                line = self.indent * indent_cur + item['instance'] + '.' + function_str
                self.code_list.append(line)
        self.code_list.append(self.indent * indent_cur + self.wait_str)

    def to_file(self, directory_path, file_path):
        if not os.path.exists(directory_path):
            os.mkdir(directory_path)
        with open(file_path, 'w') as file:
            for code in self.code_list:
                file.write(code + '\n')

    def py_to_file(self, code_str, directory_path, file_path):
        if not os.path.exists(directory_path):
            os.mkdir(directory_path)
        with open(file_path, 'w') as file:
            file.write(code_str)
