#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动生成 gRPC Python 代码脚本（相对导入版）
递归扫描 protos/ 目录并生成到 generated/
"""

import os
import subprocess
from pathlib import Path


def generate_grpc_code(proto_dir: str, output_dir: str):
    proto_path = Path(proto_dir)
    if not proto_path.exists():
        raise FileNotFoundError(f"❌ proto 目录不存在: {proto_dir}")

    proto_files = [str(p) for p in proto_path.rglob("*.proto")]
    if not proto_files:
        print("⚠️ 未找到任何 .proto 文件")
        return

    os.makedirs(output_dir, exist_ok=True)
    print(f"🔧 生成目录: {output_dir}")

    for proto_file in proto_files:
        print(f"  ⏳ 编译 {proto_file}")

        # 只需要指定根目录
        include_paths = [f"-I{proto_dir}"]

        cmd = [
            "python",
            "-m",
            "grpc_tools.protoc",
            *include_paths,
            f"--python_out={output_dir}",
            f"--grpc_python_out={output_dir}",
            proto_file,
        ]

        try:
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            if result.stdout.strip():
                print(result.stdout)
            if result.stderr.strip():
                print("⚠️ 编译警告:", result.stderr)
        except subprocess.CalledProcessError as e:
            print("❌ 编译失败:")
            print(e.stdout)
            print(e.stderr)
            continue

    # 确保每个目录都有 __init__.py
    ensure_init_files(output_dir)

    # 修改 _pb2_grpc.py 为相对导入
    convert_grpc_imports_to_relative(output_dir)

    print(f"✅ 所有 proto 文件已生成到: {output_dir}")


def ensure_init_files(output_dir: str):
    """递归生成 __init__.py 确保 Python 包可导入"""
    for root, dirs, files in os.walk(output_dir):
        init_file = os.path.join(root, "__init__.py")
        if not os.path.exists(init_file):
            open(init_file, "w", encoding="utf-8").close()


def convert_grpc_imports_to_relative(output_dir: str):
    """把 _pb2_grpc.py 中的绝对导入改为相对导入"""
    for root, dirs, files in os.walk(output_dir):
        for file in files:
            if file.endswith("_pb2_grpc.py"):
                path = os.path.join(root, file)
                with open(path, "r", encoding="utf-8") as f:
                    lines = f.readlines()
                new_lines = []
                for line in lines:
                    # 将类似 `import xxx_pb2 as xxx__pb2` 改为 `from . import xxx_pb2 as xxx__pb2`
                    if line.startswith("import ") and "_pb2 as " in line:
                        parts = line.split(" as ")
                        module_name = parts[0].replace("import", "").strip()
                        alias = parts[1].strip()
                        new_lines.append(f"from . import {module_name} as {alias}\n")
                    else:
                        new_lines.append(line)
                with open(path, "w", encoding="utf-8") as f:
                    f.writelines(new_lines)


if __name__ == "__main__":
    ROOT_DIR = Path(__file__).resolve().parents[1]
    PROTO_DIR = ROOT_DIR / "protos"
    OUTPUT_DIR = ROOT_DIR / "generated"

    print("🚀 开始生成 gRPC Python 代码（相对导入版）...")
    generate_grpc_code(str(PROTO_DIR), str(OUTPUT_DIR))
    print("\n🎉 gRPC Python 代码生成完成！")