#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UE头文件解析工具
用于扫描UE头文件并生成slua插件的emmy-lua注解
"""

import os
import re
import argparse
from pathlib import Path
from typing import List, Dict, Set, Optional

class UEHeaderParser:
    def __init__(self):
        self.classes = []
        self.structs = []
        self.enums = []
        self.delegates = []
        
    def parse_header_file(self, file_path: str) -> Dict:
        """解析单个头文件"""
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        
        result = {
            'classes': [],
            'structs': [],
            'enums': [],
            'delegates': [],
            'file_path': file_path
        }
        
        # 解析UCLASS
        class_pattern = r'UCLASS\s*\([^)]*\)\s*\n\s*class\s+[A-Z_]+\s+(\w+)\s*:\s*public\s+(\w+)'
        classes = re.findall(class_pattern, content)
        for class_name, parent_class in classes:
            result['classes'].append({
                'name': class_name,
                'parent': parent_class,
                'functions': self._parse_functions(content, class_name),
                'properties': self._parse_properties(content, class_name)
            })
        
        # 解析USTRUCT
        struct_pattern = r'USTRUCT\s*\([^)]*\)\s*\n\s*struct\s+[A-Z_]+\s+(\w+)'
        structs = re.findall(struct_pattern, content)
        for struct_name in structs:
            result['structs'].append({
                'name': struct_name,
                'properties': self._parse_struct_properties(content, struct_name)
            })
        
        # 解析USTRUCT (BlueprintType变体)
        struct_pattern2 = r'USTRUCT\s*\([^)]*\)\s*\n\s*struct\s+F(\w+)'
        structs2 = re.findall(struct_pattern2, content)
        for struct_name in structs2:
            result['structs'].append({
                'name': f'F{struct_name}',
                'properties': self._parse_struct_properties(content, f'F{struct_name}')
            })
        
        # 解析UENUM
        enum_pattern = r'UENUM\s*\([^)]*\)\s*\n\s*enum\s+class\s+(\w+)'
        enums = re.findall(enum_pattern, content)
        for enum_name in enums:
            result['enums'].append({
                'name': enum_name,
                'values': self._parse_enum_values(content, enum_name)
            })
        
        # 解析委托
        delegate_pattern = r'DECLARE_DYNAMIC_(MULTICAST_)?DELEGATE(?:_\w+)?\s*\(\s*(\w+)\s*'
        delegates = re.findall(delegate_pattern, content)
        for is_multicast, delegate_name in delegates:
            result['delegates'].append({
                'name': delegate_name,
                'is_multicast': bool(is_multicast),
                'params': self._parse_delegate_params(content, delegate_name)
            })
        
        return result
    
    def _parse_functions(self, content: str, class_name: str) -> List[Dict]:
        """解析UFUNCTION"""
        functions = []
        
        # 匹配UFUNCTION声明
        func_pattern = r'UFUNCTION\s*\([^)]*\)\s*\n\s*(\w+)\s+(\w+)\s*\(([^)]*)\)'
        matches = re.findall(func_pattern, content)
        
        for return_type, func_name, params in matches:
            # 解析参数
            param_list = []
            if params.strip():
                for param in params.split(','):
                    param = param.strip()
                    if param:
                        # 简单的参数解析
                        parts = param.split()
                        if len(parts) >= 2:
                            param_type = parts[-2] if len(parts) > 2 else parts[0]
                            param_name = parts[-1]
                            param_list.append({
                                'type': param_type,
                                'name': param_name
                            })
            
            functions.append({
                'name': func_name,
                'return_type': return_type,
                'params': param_list
            })
        
        return functions
    
    def _parse_properties(self, content: str, class_name: str) -> List[Dict]:
        """解析UPROPERTY"""
        properties = []
        
        # 匹配UPROPERTY声明
        prop_pattern = r'UPROPERTY\s*\([^)]*\)\s*\n\s*(\w+(?:<[^>]*>)?)\s+(\w+);'
        matches = re.findall(prop_pattern, content)
        
        for prop_type, prop_name in matches:
            properties.append({
                'name': prop_name,
                'type': prop_type
            })
        
        return properties
    
    def _parse_struct_properties(self, content: str, struct_name: str) -> List[Dict]:
        """解析结构体属性"""
        properties = []
        
        # 在结构体定义范围内查找属性
        struct_start = content.find(f'struct {struct_name}')
        if struct_start == -1:
            return properties
        
        # 找到结构体结束位置
        brace_count = 0
        struct_content = ""
        for i in range(struct_start, len(content)):
            char = content[i]
            if char == '{':
                brace_count += 1
            elif char == '}':
                brace_count -= 1
                if brace_count == 0:
                    struct_content = content[struct_start:i+1]
                    break
        
        if not struct_content:
            return properties
        
        # 在结构体内容中查找UPROPERTY
        prop_pattern = r'UPROPERTY\s*\([^)]*\)\s*\n\s*(\w+(?:<[^>]*>)?)\s+(\w+);'
        matches = re.findall(prop_pattern, struct_content)
        
        for prop_type, prop_name in matches:
            properties.append({
                'name': prop_name,
                'type': prop_type
            })
        
        return properties
    
    def _parse_enum_values(self, content: str, enum_name: str) -> List[Dict]:
        """解析枚举值"""
        values = []
        
        # 找到枚举定义
        enum_start = content.find(f'enum class {enum_name}')
        if enum_start == -1:
            return values
        
        # 找到枚举内容
        brace_start = content.find('{', enum_start)
        if brace_start == -1:
            return values
        
        brace_end = content.find('}', brace_start)
        if brace_end == -1:
            return values
        
        enum_content = content[brace_start+1:brace_end]
        
        # 解析枚举值
        value_pattern = r'(\w+)\s*(?:=\s*(\d+))?'
        matches = re.findall(value_pattern, enum_content)
        
        current_value = 0
        for name, explicit_value in matches:
            if explicit_value:
                current_value = int(explicit_value)
            
            values.append({
                'name': name,
                'value': current_value
            })
            
            current_value += 1
        
        return values
    
    def _parse_delegate_params(self, content: str, delegate_name: str) -> List[Dict]:
        """解析委托参数"""
        params = []
        
        # 找到委托声明
        delegate_pattern = f'DECLARE_DYNAMIC_(MULTICAST_)?DELEGATE(?:_\\w+)?\\s*\\(\\s*{delegate_name}'
        match = re.search(delegate_pattern, content)
        if not match:
            return params
        
        # 查找参数列表
        param_start = content.find('(', match.end())
        if param_start == -1:
            return params
        
        param_end = content.find(')', param_start)
        if param_end == -1:
            return params
        
        param_content = content[param_start+1:param_end]
        
        # 解析参数
        param_parts = [p.strip() for p in param_content.split(',') if p.strip()]
        for i, param in enumerate(param_parts):
            parts = param.split()
            if len(parts) >= 2:
                param_type = ' '.join(parts[:-1])
                param_name = parts[-1]
                params.append({
                    'type': param_type,
                    'name': param_name
                })
        
        return params
    
    def generate_emmy_lua_annotations(self, parsed_data: Dict) -> str:
        """生成emmy-lua注解"""
        output = []
        
        # 文件头注释
        output.append(f'-- 自动生成的emmy-lua注解文件')
        output.append(f'-- 源文件: {parsed_data["file_path"]}')
        output.append('')
        
        # 生成枚举注解
        for enum in parsed_data['enums']:
            output.append(f'---@enum {enum["name"]}')
            output.append(f'local {enum["name"]} = {{')
            for value in enum['values']:
                output.append(f'    {value["name"]} = {value["value"]},')
            output.append('}')
            output.append('')
        
        # 生成结构体注解
        for struct in parsed_data['structs']:
            output.append(f'---@class {struct["name"]}')
            for prop in struct['properties']:
                lua_type = self._cpp_to_lua_type(prop['type'])
                output.append(f'---@field {prop["name"]} {lua_type}')
            output.append(f'local {struct["name"]} = {{}}')
            output.append('')
        
        # 生成类注解
        for cls in parsed_data['classes']:
            output.append(f'---@class {cls["name"]} : {cls["parent"]}')
            
            # 添加属性
            for prop in cls['properties']:
                lua_type = self._cpp_to_lua_type(prop['type'])
                output.append(f'---@field {prop["name"]} {lua_type}')
            
            # 添加方法
            for func in cls['functions']:
                lua_return_type = self._cpp_to_lua_type(func['return_type'])
                param_annotations = []
                for param in func['params']:
                    lua_param_type = self._cpp_to_lua_type(param['type'])
                    param_annotations.append(f'---@param {param["name"]} {lua_param_type}')
                
                if param_annotations:
                    output.extend(param_annotations)
                output.append(f'---@return {lua_return_type}')
                output.append(f'function {cls["name"]}:{func["name"]}({", ".join(p["name"] for p in func["params"])}) end')
                output.append('')
            
            output.append('')
        
        # 生成委托注解
        for delegate in parsed_data['delegates']:
            output.append(f'---@class {delegate["name"]}')
            param_types = []
            for param in delegate['params']:
                lua_type = self._cpp_to_lua_type(param['type'])
                param_types.append(f'{param["name"]}: {lua_type}')
            
            if param_types:
                output.append(f'---@field Call fun({", ".join(param_types)})')
            else:
                output.append('---@field Call fun()')
            output.append(f'local {delegate["name"]} = {{}}')
            output.append('')
        
        return '\n'.join(output)
    
    def _cpp_to_lua_type(self, cpp_type: str) -> str:
        """将C++类型转换为Lua类型"""
        type_mapping = {
            'int32': 'integer',
            'int64': 'integer',
            'float': 'number',
            'double': 'number',
            'bool': 'boolean',
            'FString': 'string',
            'FText': 'string',
            'FName': 'string',
            'void': 'nil',
            'TArray': 'table',
            'TMap': 'table',
            'TSet': 'table'
        }
        
        # 处理模板类型
        if '<' in cpp_type and '>' in cpp_type:
            base_type = cpp_type.split('<')[0]
            inner_type = cpp_type.split('<')[1].split('>')[0]
            lua_inner_type = self._cpp_to_lua_type(inner_type)
            return f'{type_mapping.get(base_type, "any")}<{lua_inner_type}>'
        
        return type_mapping.get(cpp_type, 'any')
    
    def scan_directory(self, directory: str, output_dir: str = None):
        """扫描目录或单个文件并生成注解文件"""
        header_files = []
        
        if os.path.isfile(directory) and directory.endswith('.h'):
            # 单个文件
            header_files = [directory]
        elif os.path.isdir(directory):
            # 目录
            for root, dirs, files in os.walk(directory):
                for file in files:
                    if file.endswith('.h'):
                        header_files.append(os.path.join(root, file))
        else:
            print(f'错误: {directory} 不是有效的文件或目录')
            return
        
        print(f'找到 {len(header_files)} 个头文件')
        
        for header_file in header_files:
            try:
                print(f'正在解析: {header_file}')
                parsed_data = self.parse_header_file(header_file)
                
                if any([parsed_data['classes'], parsed_data['structs'], parsed_data['enums'], parsed_data['delegates']]):
                    annotations = self.generate_emmy_lua_annotations(parsed_data)
                    
                    # 确定输出文件路径
                    if output_dir:
                        if os.path.isfile(directory):
                            # 单个文件的情况
                            output_file = os.path.join(output_dir, os.path.basename(header_file).replace('.h', '.d.lua'))
                        else:
                            # 目录的情况
                            relative_path = os.path.relpath(header_file, directory)
                            output_file = os.path.join(output_dir, relative_path.replace('.h', '.d.lua'))
                        os.makedirs(os.path.dirname(output_file), exist_ok=True)
                    else:
                        output_file = header_file.replace('.h', '.d.lua')
                    
                    with open(output_file, 'w', encoding='utf-8') as f:
                        f.write(annotations)
                    
                    print(f'已生成: {output_file}')
                
            except Exception as e:
                print(f'解析文件 {header_file} 时出错: {e}')

def main():
    parser = argparse.ArgumentParser(description='UE头文件解析工具 - 生成slua插件的emmy-lua注解')
    parser.add_argument('directory', help='要扫描的目录路径')
    parser.add_argument('-o', '--output', help='输出目录路径（默认为源文件同目录）')
    parser.add_argument('--recursive', action='store_true', help='递归扫描子目录')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.directory):
        print(f'错误: 目录 {args.directory} 不存在')
        return
    
    parser = UEHeaderParser()
    parser.scan_directory(args.directory, args.output)

if __name__ == '__main__':
    main()