import os
import subprocess
from pathlib import Path

def generate_proto():
    # 获取当前目录
    current_dir = Path(__file__).parent
    proto_dir = current_dir / "protos"
    generated_dir = current_dir / "generated"

    # 创建生成的目录
    generated_dir.mkdir(parents=True, exist_ok=True)

    # 编译所有 .proto 文件
    proto_files = []
    for root, _, files in os.walk(proto_dir):
        for file in files:
            if file.endswith(".proto"):
                proto_files.append(os.path.join(root, file))

    for proto_path in proto_files:
        # 计算相对路径
        rel_path = os.path.relpath(proto_path, proto_dir)
        output_dir = generated_dir

        cmd = [
            "python", "-m", "grpc_tools.protoc",
            f"-I{proto_dir}",
            f"--python_out={output_dir}",
            f"--grpc_python_out={output_dir}",
            proto_path
        ]

        print(f"Generating code for {rel_path}...")
        result = subprocess.run(cmd, capture_output=True, text=True)

        if result.returncode != 0:
            print(f"Error generating {rel_path}:")
            print(result.stderr)
            return False

    print("Proto files generated successfully!")

    # 修复生成的代码中的导入路径
    fix_imports(generated_dir)

    return True

def fix_imports(generated_dir):
    """修复生成的代码中的导入路径"""
    for root, _, files in os.walk(generated_dir):
        for file in files:
            if file.endswith(".py"):
                file_path = os.path.join(root, file)
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()

                # 替换导入路径
                content = content.replace(
                    "from cmvr.api import",
                    "from generated.cmvr.api import"
                )
                content = content.replace(
                    "import cmvr.api.",
                    "import generated.cmvr.api."
                )

                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(content)

                print(f"Fixed imports in {file_path}")

if __name__ == "__main__":
    generate_proto()