import os
import zipfile
from io import BytesIO
from pathlib import Path
from mako.lookup import TemplateLookup
import black

class CodeGenerator:
    def __init__(self, template_root='templates'):
        self.template_root = template_root
        self.lookup = TemplateLookup(directories=[template_root], input_encoding='utf-8')
    
    def generate_files(self, entity_config, prj_template_dir, params, language='python'):
        """
        生成所有匹配的模板文件
        :param entity_config: 实体配置
        :param prj_template_dir: 项目对应的文件路径
        :param params: 生成参数
        :return: 生成的文件字典 {相对路径: 文件内容}
        """
        entity_config = type("EntityConfig", (object,), entity_config)
        generated_files = []
        template_dir = os.path.join(self.template_root, prj_template_dir)
        
        # 遍历模板目录
        for root, _, files in os.walk(template_dir):
            for file in files:
                if not file.endswith('.mako'): continue
                template_path = os.path.join(root, file)
                relative_path = os.path.relpath(template_path, self.template_root)
                
                # 渲染模板
                template = self.lookup.get_template(relative_path)
                content = template.render(entity=entity_config, params=params, generator=self)

                if relative_path.endswith('.py.mako'):
                    content = self.post_process(content)
                
                # 计算输出路径
                output_path = self._get_output_path(relative_path, entity_config, params)
                # generated_files[output_path] = content
                generated_files.append(dict(
                    path=output_path,
                    content=content
                ))
        
        return generated_files
    
    def post_process(self, content):
        try:
            # 尝试使用 Black 格式化代码
            formatted_code = black.format_str(content, mode=black.Mode(line_length=100))
            return formatted_code
        except Exception as e:
            print(f"Black 格式化失败: {e}")
            return content
    
    def _get_output_path(self, template_path, entity_config, params):
        """
        根据模板路径计算输出路径
        :param template_path: 模板相对路径
        :return: 输出文件相对路径(无.mako后缀)
        """
        # 移除.mako后缀
        if template_path.endswith('.mako'):
            output_path = template_path[:-5]
        else:
            output_path = template_path
        
        # 替换实体名称占位符
        entity_name = entity_config.name
        output_path = output_path.replace('Entity', entity_name.capitalize())
        output_path = output_path.replace('entity', entity_name.lower())
        
        return output_path
    
    def _process_relations(self, entity):
        """处理实体关系并返回处理后的结果"""
        if not hasattr(entity, 'relations') or not entity.relations:
            return []
            
        processed = []
        for rel in entity.relations:
            # 确保关系有必要的字段
            rel.setdefault('through', None)
            rel.setdefault('description', '')
            
            # 计算反向引用名称
            rel['reverse_name'] = self._get_reverse_relation_name(entity.name, rel)
            processed.append(rel)
        return processed
    
    def _get_reverse_relation_name(self, entity_name, relation):
        """生成反向引用名称"""
        if relation['type'] == 'one-to-many':
            return entity_name.lower()
        elif relation['type'] == 'many-to-one':
            return relation['name']
        elif relation['type'] == 'many-to-many':
            return f"{entity_name.lower()}s"
        return None
    
    def create_zip(self, generated_files):
        """
        将生成的文件打包为zip
        :param generated_files: {路径: 内容} 字典
        :return: zip文件字节流
        """
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
           for item in generated_files:
                path, content = item.get('path'), item.get('content')
                zip_file.writestr(path, content)
        zip_buffer.seek(0)
        return zip_buffer
    
    def create_local_code(self, generated_files, output_dir):
        save_dir = Path(output_dir)

        for item in generated_files:
            # 确保输出目录存在
            path, content = item.get('path'), item.get('content')
            output_file = save_dir / path
            output_file.parent.mkdir(parents=True, exist_ok=True)
           
            with open(output_file, 'w', encoding='utf-8', newline='') as f:
                f.write(content)


if __name__ == "__main__":
    # 使用示例
    generator = CodeGenerator()

    default_frontend_config = {
        "editable": True,
        "listable": True,
        "detailable": True,
        "sortable": True,
        "filterable": True,
        "filter_operator": 'equal',  # equal | contains 默认过滤操作符为包含
        "display_type": 'text',  # 默认显示类型为文本
    }
   
    entity_config = {
        "name": "user",  # 实体名称（小写）
        "author": "Ly997",  # 实体名称（小写）
        "description": "系统用户管理",  # 实体描述
        "version": "1.0.0",  # 版本号
        "fields": [
            {
                "name": "id",
                "type": "Long",
                "primary_key": True,
                "description": "用户ID",
                "required": True
            },
            {
                "name": "username",
                "type": "String",
                "length": 50,
                "description": "用户名",
                "required": True,
                "unique": True,
                "validations": [
                    {"type": "min_length", "value": 4},
                    {"type": "max_length", "value": 20}
                ]
            },
            {
                "name": "email",
                "type": "String",
                "description": "电子邮箱",
                "required": True,
                "validations": [
                    {"type": "email"}
                ]
            },
            {
                "name": "password_hash",
                "type": "String",
                "description": "密码哈希",
                "required": True,
                "secret": True  # 标记为敏感字段
            },
            {
                "name": "is_active",
                "type": "Boolean",
                "description": "是否激活",
                "default": True
            },
            {
                "name": "created_at",
                "type": "DateTime",
                "description": "创建时间",
                "auto_now_add": True
            },
            {
                "name": "updated_at",
                "type": "DateTime",
                "description": "更新时间",
                "auto_now": True
            }
        ],
        "relations": [
            {
                "name": "roles",
                "type": "many-to-many",
                "target": "role",
                "through": "user_roles",
                "description": "用户角色关联"
            },
            {
                "name": "posts",
                "type": "one-to-many",
                "target": "post",
                "description": "用户发表的文章"
            }
        ],
        "api": {
            "operations": ["create", "read", "update", "delete", "list"],
            "base_path": "/users",
            "auth_required": True,
            "permissions": {
                "create": ["admin", "manager"],
                "delete": ["admin"]
            }
        },
        "menu": {
            "label": "用户管理",
            "icon": "user",
            "order": 10,
            "submenu": [
                {
                    "label": "用户列表",
                    "path": "/users",
                    "icon": "list"
                },
                {
                    "label": "角色管理",
                    "path": "/roles",        
                }
            ]
        }
    }
    params = {
        "package": "com.example",
        "author": "John Doe"
    }

    # 在代码中合并默认配置和字段的前端配置
    for field in entity_config["fields"]:
        frontend = {**default_frontend_config, **field}
        field.update(frontend)

    import json
    with open('tmp.json', 'w', encoding='utf-8') as f:
        json.dump(entity_config, f, ensure_ascii=False, indent=4)

    # 生成所有文件
    generated_files = generator.generate_files(entity_config, 'relation-demo', params)
    # generated_files = generator.generate_files(entity_config, 'vue-fastapi-admin', params)
    generator.create_local_code(generated_files, 'tmp')
    # 打包为zip
    zip_data = generator.create_zip(generated_files)

    file_path = './tmp.zip'
    with open(file_path, 'wb') as f:
        f.write(zip_data.getvalue())
        print(f"文件已保存到 {file_path}")

    # 在Flask中返回zip文件示例
    # return send_file(zip_data, mimetype='application/zip', as_attachment=True, download_name='generated_code.zip')
