# TransPyC 入口程序

import sys
import os
import ast
import json
import struct
# 添加lib目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'lib', 'includes'))
from lib.includes import c
from lib.includes import t
from lib.constants.config import (
    DEFAULT_INPUT_FILE, DEFAULT_OUTPUT_FILE,
    DEFAULT_COMPILE_COMMAND,
    DEFAULT_COMPILE_FLAGS, ERROR_MESSAGES,
    HELP_MESSAGE, DEFAULT_METADATA
)
from lib.utils.helpers import (
    DetectFileType, ExecuteCommand,
    AppendFileContent
)
from lib.core.translator import Translator


def SerializeSymbolTable(SymbolTable: dict) -> bytes:
    """将符号表序列化为二进制格式

    格式:
    - 4字节: JSON数据长度（小端序）
    - N字节: JSON数据（UTF-8编码）

    Args:
        SymbolTable: 符号表字典

    Returns:
        二进制字节数据
    """
    JsonData = json.dumps(SymbolTable, ensure_ascii=False).encode('utf-8')
    length = len(JsonData)
    # 使用4字节小端序整数表示长度
    header = struct.pack('<I', length)
    return header + JsonData


def deSerializeSymbolTable(data: bytes) -> dict:
    """从二进制数据反序列化符号表

    Args:
        data: 二进制字节数据

    Returns:
        符号表字典
    """
    if len(data) < 4:
        raise ValueError("Invalid symbin file: data too short")
    # 解析4字节长度头（小端序）
    length = struct.unpack('<I', data[:4])[0]
    JsonData = data[4:4+length]
    return json.loads(JsonData.decode('utf-8'))


class SymbolFile:
    """符号文件类，用于解析和存储符号信息"""

    def __init__(self, file=None, string=None, type=None, encoding='utf-8'):
        """初始化SymbolFile对象

        Args:
            file: 文件路径
            string: 代码字符串
            type: 文件类型，支持"c"或"py"
            encoding: 文件编码
        """
        self.FilePath = file
        self.CodeString = string
        self.FileType = type
        self.encoding = encoding
        self.symbols = {}



class Config:
    """配置类"""
    def __init__(self):
        self.debug = False
        self.mode = 'relaxed'  # 'relaxed' 或 'strict'


class TransPyC:
    """TransPyC 主类"""

    def __init__(self, code=None, debug=False, triple=None, datalayout=None):
        """初始化TransPyC对象

        Args:
            code: 代码字符串
            debug: 调试模式
            triple: LLVM target triple (e.g. "x86_64-pc-linux-gnu")
            datalayout: LLVM data layout string
        """
        self.Args = {}
        self.HeaderFiles = []
        self.HelperFiles = []
        self.AnnotationFiles = []
        self.translator = Translator()
        self.code = code
        self.translator.Content = code
        self.config = Config()
        self.translator.config = self.config
        self.config.debug = debug
        self.SymbolFiles = []
        self.SliceLevel = 3
        self.SliceCount = 0
        self.SliceInfos = []
        self.metadata = DEFAULT_METADATA.copy()
        self.triple = triple
        self.datalayout = datalayout
        self._source_module_sig_files = {}

    @staticmethod
    def PreProcessSymbol(SymbolFile, debug=False):
        """预处理符号文件，提取符号表并返回二进制数据

        Args:
            SymbolFile: SymbolFile对象，包含文件路径和类型信息
            debug: 是否启用调试模式，如果为True，返回 (bytes, DebugInfo)
                  对于Python文件，DebugInfo包含p2c日志
                  对于C文件，DebugInfo包含基本处理信息

        Returns:
            如果debug为False: bytes - 序列化后的符号表二进制数据
            如果debug为True: (bytes, str) - (二进制数据, 调试日志)
        """
        translator = Translator()
        DebugLogs = []

        try:
            # 读取文件内容
            if SymbolFile.FilePath:
                with open(SymbolFile.FilePath, 'r', encoding=SymbolFile.encoding) as f:
                    content = f.read()
            elif SymbolFile.CodeString:
                content = SymbolFile.CodeString
            else:
                raise ValueError("SymbolFile must have either FilePath or CodeString")

            # 根据文件类型处理
            FileType = SymbolFile.FileType
            if not FileType and SymbolFile.FilePath:
                # 从文件路径推断类型
                if SymbolFile.FilePath.endswith('.py'):
                    FileType = 'py'
                elif SymbolFile.FilePath.endswith('.c') or SymbolFile.FilePath.endswith('.h'):
                    FileType = 'c'

            if FileType == 'py':
                # Python文件处理
                translator.OriginalLines = content.split('\n')
                translator.Content = content

                # 设置调试文件（如果需要）
                if debug:
                    DebugFile = SymbolFile.FilePath.replace('.py', '.p2c') if SymbolFile.FilePath else 'debug.p2c'
                    translator.SetDebugFile(DebugFile)
                    # 清空调试文件
                    with open(DebugFile, 'w', encoding=SymbolFile.encoding) as f:
                        f.write('')

                # 解析AST并提取符号
                tree = ast.parse(content)

                if debug:
                    with open(DebugFile, 'a', encoding=SymbolFile.encoding) as f:
                        f.write('=== AST Tree (Compact) ===\n')
                        f.write(ast.dump(tree))
                        f.write('\n\n')

                # 遍历AST提取符号
                for node in ast.iter_child_nodes(tree):
                    if isinstance(node, ast.ClassDef):
                        translator.SymbolTable[node.name] = {'type': 'struct', 'members': {}}
                        for item in node.body:
                            if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
                                VarName = item.target.id
                                TypeInfo = translator.TypeMergeHandler.GetCTypeInfo(item.annotation) if item.annotation else None
                                VarType = TypeInfo.ToString() if TypeInfo else 'unknown'
                                IsPtr = TypeInfo.IsPtr if TypeInfo else False
                                translator.SymbolTable[node.name]['members'][VarName] = {
                                    'type': VarType,
                                    'IsPtr': IsPtr
                                }
                    elif isinstance(node, ast.FunctionDef):
                        translator.SymbolTable[node.name] = {'type': 'function'}
                    elif isinstance(node, ast.AnnAssign):
                        if isinstance(node.target, ast.Name):
                            VarName = node.target.id
                            VarType = 'unknown'
                            IsPtr = False
                            if node.annotation:
                                try:
                                    TypeInfo = translator.TypeMergeHandler.GetCTypeInfo(node.annotation)
                                    VarType = TypeInfo.ToString() if TypeInfo else 'unknown'
                                    IsPtr = TypeInfo.IsPtr if TypeInfo else False
                                except:
                                    pass
                            translator.SymbolTable[VarName] = {
                                'type': 'variable',
                                'declared_type': VarType,
                                'IsPtr': IsPtr
                            }

                if debug:
                    with open(DebugFile, 'a', encoding=SymbolFile.encoding) as f:
                        f.write('=== Symbol Table ===\n')
                        f.write(str(translator.SymbolTable))
                        f.write('\n\n')
                    # 读取调试日志
                    with open(DebugFile, 'r', encoding=SymbolFile.encoding) as f:
                        DebugLogs = f.read()

            elif FileType == 'c':
                # C文件处理 - 使用现有的ParseHelperFiles逻辑
                translator.ParseHelperFiles([SymbolFile.FilePath], SymbolFile.encoding)

                if debug:
                    DebugLogs = f"=== C File Symbol Extraction ===\n"
                    DebugLogs += f"File: {SymbolFile.FilePath}\n"
                    DebugLogs += f"=== Symbol Table ===\n"
                    DebugLogs += str(translator.SymbolTable)
                    DebugLogs += "\n\n"
            else:
                raise ValueError(f"Unsupported file type: {FileType}")

            # 序列化符号表
            BinaryData = SerializeSymbolTable(translator.SymbolTable)

            if debug:
                return BinaryData, DebugLogs
            return BinaryData

        except Exception as e:
            error_msg = f"Error processing symbol file: {e}"
            if debug:
                return b'', error_msg
            raise RuntimeError(error_msg)

    def ParseArgs(self):
        """解析命令行参数"""
        I = 1
        while I < len(sys.argv):
            if sys.argv[I] == '-f':
                if I + 1 < len(sys.argv):
                    self.Args['Input'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -f requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-o':
                if I + 1 < len(sys.argv):
                    self.Args['Output'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -o requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-wh':
                if I + 1 < len(sys.argv):
                    self.HeaderFiles = sys.argv[I + 1:]
                    break
                else:
                    print(f'Error: -wh requires arguments')
                    sys.exit(1)
            elif sys.argv[I] == '-debug':
                if I + 1 < len(sys.argv):
                    self.Args['Debug'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -debug requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-cc':
                # 编译命令
                if I + 1 < len(sys.argv):
                    self.Args['CompileCommand'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -cc requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-cflags':
                # 编译标志
                if I + 1 < len(sys.argv):
                    self.Args['CompileFlags'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -cflags requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-run':
                # 是否运行生成的程序
                self.Args['Run'] = True
                I += 1
            elif sys.argv[I] == '-args':
                # 运行时参数
                if I + 1 < len(sys.argv):
                    self.Args['RunArgs'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -args requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-level' or sys.argv[I] == '-l':
                if I + 1 < len(sys.argv):
                    LevelArg = sys.argv[I + 1]
                    if LevelArg.isdigit():
                        self.Args['SliceLevel'] = int(LevelArg)
                    elif LevelArg == 's':
                        self.Args['SliceLevel'] = 's'
                    else:
                        self.Args['SliceLevel'] = 3
                    I += 1
                else:
                    self.Args['SliceLevel'] = 3
            elif sys.argv[I] == '-h':
                # 辅助文件，用于解析符号信息
                if I + 1 < len(sys.argv):
                    # 收集所有后续参数作为辅助文件，直到遇到下一个以-开头的参数
                    while I + 1 < len(sys.argv) and not sys.argv[I + 1].startswith('-'):
                        self.HelperFiles.append(sys.argv[I + 1])
                        I += 1
                    I += 1
                else:
                    print(f'Error: -h requires arguments')
                    sys.exit(1)
            elif sys.argv[I] == '-presym':
                # 预处理符号文件，生成.symbin文件
                # 格式: -presym <InputFile> -o <OutputFile> [-debug <DebugFile>]
                if I + 1 < len(sys.argv):
                    InputFile = sys.argv[I + 1]
                    I += 2

                    # 解析可选参数
                    OutputFile = None
                    DebugFile = None
                    while I < len(sys.argv) and sys.argv[I].startswith('-'):
                        if sys.argv[I] == '-o':
                            if I + 1 < len(sys.argv):
                                OutputFile = sys.argv[I + 1]
                                I += 2
                            else:
                                print(f'Error: -o requires an argument')
                                sys.exit(1)
                        elif sys.argv[I] == '-debug':
                            if I + 1 < len(sys.argv):
                                DebugFile = sys.argv[I + 1]
                                I += 2
                            else:
                                print(f'Error: -debug requires an argument')
                                sys.exit(1)
                        else:
                            break

                    # 如果没有指定输出文件，使用默认名称
                    if not OutputFile:
                        OutputFile = InputFile.replace('.py', '.symbin').replace('.c', '.symbin')

                    # 执行预处理
                    self.Args['PreSym'] = {
                        'input': InputFile,
                        'output': OutputFile,
                        'debug': DebugFile
                    }
                else:
                    print(f'Error: -presym requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '-e' or sys.argv[I] == '--encoding':
                if I + 1 < len(sys.argv):
                    self.Args['Encoding'] = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: -e/--encoding requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '--triple':
                if I + 1 < len(sys.argv):
                    self.triple = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: --triple requires an argument')
                    sys.exit(1)
            elif sys.argv[I] == '--datalayout':
                if I + 1 < len(sys.argv):
                    self.datalayout = sys.argv[I + 1]
                    I += 2
                else:
                    print(f'Error: --datalayout requires an argument')
                    sys.exit(1)
            else:
                print(f'Error: Unknown argument {sys.argv[I]}')
                sys.exit(1)

        # 检查是否有预处理符号文件的请求，如果有则跳过输入输出检查
        if 'PreSym' not in self.Args and ('Input' not in self.Args or 'Output' not in self.Args):
            print(ERROR_MESSAGES['MISSING_ARGS'])
            print(HELP_MESSAGE)
            sys.exit(1)

    def CompileAndRun(self, OutputFile):
        """编译并运行生成的C代码"""
        # 获取编译命令
        CompileCmd = self.Args.get('CompileCommand', DEFAULT_COMPILE_COMMAND)
        CompileFlags = self.Args.get('CompileFlags', DEFAULT_COMPILE_FLAGS)

        # 构建编译命令
        FullCompileCmd = f'{CompileCmd} {CompileFlags} {OutputFile} -o {OutputFile}.exe'
        print(f'Compiling: {FullCompileCmd}')

        # 执行编译命令
        try:
            CompileResult = ExecuteCommand(FullCompileCmd)
            if CompileResult:
                print('Compilation successful!')

                # 检查是否需要运行
                if self.Args.get('Run', False):
                    RunArgs = self.Args.get('RunArgs', '')
                    RunCmd = f'{OutputFile}.exe {RunArgs}'.strip()
                    print(f'Running: {RunCmd}')

                    # 执行运行命令
                    RunResult = ExecuteCommand(RunCmd)
                    if RunResult:
                        print('Execution output:')
                        print(RunResult.stdout)
                        if RunResult.stderr:
                            print('Execution errors:')
                            print(RunResult.stderr)
        except Exception as e:
            print(f'{ERROR_MESSAGES["COMPILE_FAILED"]}: {e}')

    def WriteDebugInfo(self, content):
        """写入调试信息到指定文件"""
        # 获取调试文件路径
        if 'Debug' in self.Args:
            DebugFile = self.Args['Debug']
        else:
            DebugFile = self.Args.get('Output', '').replace('.c', '.p2c')

        if DebugFile:
            AppendFileContent(DebugFile, content)

    def PythonToC(self, InputFile, OutputFile):
        """Python到C的转换"""
        # 获取编码参数
        encoding = self.Args.get('Encoding', 'utf-8')

        # 设置调试文件路径（使用 .p2c 扩展名）
        if 'Debug' in self.Args:
            DebugFile = self.Args['Debug']
        else:
            # 默认使用 .p2c 文件
            DebugFile = OutputFile.replace('.c', '.p2c')

        # 设置翻译器的调试文件
        self.translator.SetDebugFile(DebugFile)

        # 先清空调试文件
        with open(DebugFile, 'w', encoding=encoding) as f:
            f.write('')

        # 解析辅助文件，提取符号信息
        if self.HelperFiles:
            self.translator.ParseHelperFiles(self.HelperFiles, encoding)

        # 加载注解文件，提取类型定义
        if self.AnnotationFiles:
            self.translator.LoadAnnotationFiles(self.AnnotationFiles)

        # 获取编码参数
        encoding = self.Args.get('Encoding', 'utf-8')
        with open(InputFile, 'r', encoding=encoding) as F:
            Content = F.read()

        # 保存原始代码行
        self.translator.OriginalLines = Content.split('\n')
        self.translator.Content = Content

        # 解析Python代码为AST
        Tree = ast.parse(Content)

        # 解析文件以提取类型信息（用于 EmbeddedAssignments）
        self.translator.ParsePythonFile(InputFile, encoding)

        # 写入AST树信息（压缩格式）
        with open(DebugFile, 'a', encoding=encoding) as f:
            f.write('=== AST Tree (Compact) ===\n')
            f.write(ast.dump(Tree))
            f.write('\n\n')

        Target = 'llvm'
        CCode = self.translator.GenerateCCode(Tree, target=Target)

        import datetime
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        filename = os.path.basename(OutputFile)

        encoding = self.Args.get('Encoding', 'utf-8')
        with open(OutputFile, 'w', encoding=encoding) as F:
            F.write(CCode)
            F.write('\n')

    def CToPython(self, InputFile, OutputFile, HeaderFiles):
        """C到Python的转换（暂时禁用）"""
        print("C到Python的转换功能暂时禁用，请使用Python到C的转换功能。")

    def AddSymbol(self, SymbolFiles):
        """添加符号文件

        Args:
            SymbolFiles: 符号文件或符号文件列表
        """
        if isinstance(SymbolFiles, list):
            self.SymbolFiles.extend(SymbolFiles)
        else:
            self.SymbolFiles.append(SymbolFiles)

        # 收集文件路径和编码
        HelperFiles = []
        encodings = []
        for SymbolFile in self.SymbolFiles:
            FilePath = SymbolFile.FilePath
            encoding = SymbolFile.encoding
            if FilePath:
                HelperFiles.append(FilePath)
                encodings.append(encoding)

        # 解析辅助文件
        if HelperFiles:
            # 暂时使用第一个文件的编码
            encoding = encodings[0] if encodings else 'utf-8'
            self.translator.ParseHelperFiles(HelperFiles, encoding)

    def Convert(self, OutputFilename=None, SourceFilename=None, target='llvm'):
        """转换代码

        Args:
            OutputFilename: 输出文件名（用于模板渲染）
            SourceFilename: 源文件路径（用于设置CurrentFile，影响导入模块的查找）
            target: 目标代码类型，仅支持 'llvm'

        Returns:
            如果debug是字符串，则返回生成的代码
            如果debug是True，则返回生成的代码和调试信息
        """
        if not self.code:
            print('Error: No code provided')
            return None

        from lib.core.Handles.HandlesImports import ImportHandle
        ImportHandle._struct_load_cache.clear()

        # 如果提供了 SourceFilename，先设置 CurrentFile
        # 这样导入模块时会相对于源文件所在目录查找
        if SourceFilename:
            self.translator.CurrentFile = SourceFilename

        # 加载注解文件，提取类型定义
        if self.AnnotationFiles:
            self.translator.LoadAnnotationFiles(self.AnnotationFiles)

        # 保存原始代码行
        self.translator.OriginalLines = self.code.split('\n')

        # 解析Python代码为AST
        Tree = ast.parse(self.code)

        # 从源代码直接预扫描 import 语句，注册别名
        # 注意：不能依赖 Tree 中的 import 节点，因为 ParsePythonFile 可能修改了 AST
        if not getattr(self.translator, '_import_aliases', None):
            self.translator._import_aliases = {}
        if not getattr(self.translator, '_imported_modules', None):
            self.translator._imported_modules = set()
        _prescan_tree = ast.parse(self.code)
        for _node in ast.iter_child_nodes(_prescan_tree):
            if isinstance(_node, ast.Import):
                for _alias in _node.names:
                    _name = _alias.name
                    if _name in ('c', 't'):
                        continue
                    self.translator._imported_modules.add(_name)
                    if _alias.asname:
                        self.translator._import_aliases[_alias.asname] = _name
            elif isinstance(_node, ast.ImportFrom):
                _module_name = _node.module
                if _module_name and _module_name not in ('c', 't'):
                    self.translator._imported_modules.add(_module_name)
                    for _alias in _node.names:
                        if _alias.asname:
                            self.translator._import_aliases[_alias.asname] = f"{_module_name}.{_alias.name}"

        # 解析代码以提取类型信息（用于 EmbeddedAssignments 和 TypedefAssignments）
        # 创建临时文件路径用于解析
        import tempfile
        import os
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False, encoding='utf-8') as f:
            f.write(self.code)
            TempFile = f.name
        try:
            # UpdateCurrentFile=False 表示不修改 CurrentFile
            if SourceFilename:
                self.translator.ParsePythonFile(TempFile, UpdateCurrentFile=False)
            else:
                self.translator.ParsePythonFile(TempFile)
        finally:
            os.unlink(TempFile)

        import datetime
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')

        if self.triple:
            self.translator.triple = self.triple
        if self.datalayout:
            self.translator.datalayout = self.datalayout

        # 生成代码
        code = self.translator.GenerateCCode(Tree, target='llvm')
        filename = os.path.basename(OutputFilename) if OutputFilename else 'generated.ll'

        # 处理调试信息
        if isinstance(self.config.debug, str):
            # 如果debug是字符串，写入调试信息到文件
            with open(self.config.debug, 'w', encoding='utf-8') as f:
                f.write('=== Symbol Table ===\n')
                f.write(str(self.translator.SymbolTable) + '\n\n')
                f.write('=== AST Tree ===\n')
                f.write(ast.dump(Tree, indent=2) + '\n\n')
                f.write('=== Export Table ===\n')
                f.write(self.translator.ExportTable.to_json() + '\n\n')
                f.write('=== Generated Code ===\n')
                f.write(code + '\n')
            return code
        elif self.config.debug:
            # 如果debug是True，返回生成的代码和调试信息
            DebugInfo = {
                'SymbolTable': self.translator.SymbolTable,
                'ast_tree': ast.dump(Tree, indent=2),
                'export_table': self.translator.ExportTable.to_dict(),
                'generated_code': code
            }
            return code, DebugInfo
        else:
            # 如果debug是False，只返回生成的代码
            return code

    def Run(self):
        """主运行函数"""
        # 检查是否有命令行参数
        if len(sys.argv) > 1:
            # 有命令行参数，使用 ParseArgs 方法解析
            self.ParseArgs()

            # 检查是否有预处理符号文件的请求
            if 'PreSym' in self.Args:
                PresymConfig = self.Args['PreSym']
                InputFile = PresymConfig['input']
                OutputFile = PresymConfig['output']
                DebugFile = PresymConfig['debug']

                # 获取编码参数
                encoding = self.Args.get('Encoding', 'utf-8')

                # 推断文件类型
                FileType = 'py' if InputFile.endswith('.py') else 'c' if InputFile.endswith('.c') else None

                # 创建 SymbolFile 对象，传入编码参数
                SymbolFile = SymbolFile(file=InputFile, type=FileType, encoding=encoding)

                # 调用 PreProcessSymbol
                if DebugFile:
                    BinaryData, DebugInfo = TransPyC.PreProcessSymbol(SymbolFile, debug=True)
                    # 写入调试文件
                    with open(DebugFile, 'w', encoding=encoding) as f:
                        f.write(DebugInfo)
                    print(f"Debug info written to: {DebugFile}")
                else:
                    BinaryData = TransPyC.PreProcessSymbol(SymbolFile, debug=False)

                # 写入二进制符号文件
                with open(OutputFile, 'wb') as f:
                    f.write(BinaryData)

                print(f"Symbol file generated: {OutputFile}")
                return

            InputFile = self.Args['Input']
            OutputFile = self.Args['Output']
        else:
            # 没有命令行参数，使用默认的测试文件名称
            InputFile = DEFAULT_INPUT_FILE
            OutputFile = DEFAULT_OUTPUT_FILE
            print('Using default test files: test.py -> test.c')

        FileType = DetectFileType(InputFile)

        if FileType == '.py':
            self.PythonToC(InputFile, OutputFile)
            # 检查是否需要编译和运行
            if 'CompileCommand' in self.Args or self.Args.get('Run', False):
                self.CompileAndRun(OutputFile)
        elif FileType == '.c':
            self.CToPython(InputFile, OutputFile, self.HeaderFiles)
        else:
            print(f'Error: Unsupported file type {FileType}')
            sys.exit(1)


if __name__ == '__main__':
    trans = TransPyC()
    if 'SliceLevel' in trans.Args:
        trans.SliceLevel = trans.Args['SliceLevel']
    trans.Run()