#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# proto -> Go 生成脚本（yunyan-sas monorepo）
#
# 用法：在本目录(apps/proto)下执行 `python3 pb.py`（也可从任意目录执行，脚本会自定位）。
#
# 约定：
#   - proto 源码根目录 = 本脚本所在目录(apps/proto)；所有 .proto 的 go_package 均为 ".;pb"，
#     故生成的 *.pb.go 一律平铺输出到 apps/services/pb（basename 命名，忽略子目录层级）。
#   - -I 指向 proto 根目录，用于解析形如 `import "db.proto"` / `import "api/api_db.proto"` 的引用。
#   - 递归自动发现全部 .proto（EXCLUDE_DIRS 除外），新增 proto/子目录无需改本脚本 —— 方便以后做
#     数据结构与协议更新：改完 .proto 直接重跑本脚本即可。
#   - 生成后把 .proto 里 `//@go_tags(`gorm:"..." json:"..."`)` 注释回填成真正的 struct tag
#     （protoc-gen-go 本身不写 gorm/json tag，靠本步骤注入；写死在 .pb.go 里手改无效，须改 .proto 重生成）。
#
# 前置：protoc、protoc-gen-go(v1.36.x，在 $PATH 或 $GOBIN 里)。

import io
import os
import re
import subprocess

# 本脚本所在目录 = proto 源码根；生成产物到 apps/services/pb。
PROTO_ROOT = os.path.dirname(os.path.abspath(__file__))
OUT_DIR = os.path.abspath(os.path.join(PROTO_ROOT, os.pardir, 'services', 'pb'))

# 不参与编译的目录：google/* 仅为 import 提供 well-known 类型定义，编译会与 Go 运行时内置类型冲突。
EXCLUDE_DIRS = {'google'}


def apply_go_tags(pbgo_file):
    """把 protoc 生成文件里的 //@go_tags(`k:"v" ...`) 注释回填成真正的 struct tag。

    合并规则：保留 protoc 原本写入的 tag（如 protobuf:"..."），再叠加 @go_tags 里声明的
    键值对（如 gorm/json），同名以 @go_tags 为准。
    """
    file_data = ""
    with io.open(pbgo_file, "r", encoding='utf-8') as f:
        for line in f:
            if '@go_tags' in line:
                tags = {}
                tagsstr = ""
                # 1. 解析 protoc 原有反引号 tag（第一个反引号组）
                existing = re.findall(r"`(.+?)`", line)
                if existing:
                    for v in existing[0].split():
                        tag = v.split(':')
                        if len(tag) >= 2:
                            tags[tag[0]] = tag[1]
                # 2. 解析 @go_tags(`...`) 里的 key:"value"，覆盖同名
                go_tags_content = re.findall(r"@go_tags\(\`(.+?)\`", line)
                if go_tags_content:
                    for key, value in re.findall(r'(\w+):"([^"]*)"', go_tags_content[0]):
                        tags[key] = f"\"{value}\""
                # 3. 重建 tag 串并替换回该行第一个反引号组
                for key, value in tags.items():
                    tagsstr += "{0}:{1} ".format(key, value)
                line = re.sub(r"`([^`]*)`", "`{0}`".format(tagsstr.strip()), line, count=1)
            file_data += line
    with io.open(pbgo_file, "w", encoding='utf-8') as f:
        f.write(file_data)


def build_one(proto_file):
    """编译单个 .proto 并回填 go_tags。proto_file 为绝对路径。"""
    rel = os.path.relpath(proto_file, PROTO_ROOT)
    cmd = ['protoc', '--go_out=' + OUT_DIR, '-I', PROTO_ROOT, rel]
    print(' '.join(cmd))
    if subprocess.call(cmd, cwd=PROTO_ROOT) != 0:
        raise SystemExit('protoc 编译失败: ' + rel)
    # go_package=".;pb" -> 平铺输出为 <OUT_DIR>/<basename>.pb.go
    base = os.path.splitext(os.path.basename(proto_file))[0]
    out = os.path.join(OUT_DIR, base + '.pb.go')
    if os.path.exists(out):
        apply_go_tags(out)
    else:
        print('  警告：未找到生成文件 {0}（检查 .proto 的 go_package 是否为 ".;pb"）'.format(out))


def discover_protos():
    """递归发现 PROTO_ROOT 下全部 .proto（跳过 EXCLUDE_DIRS）。"""
    protos = []
    for dirpath, dirnames, filenames in os.walk(PROTO_ROOT):
        dirnames[:] = [d for d in dirnames if d not in EXCLUDE_DIRS]
        for fn in filenames:
            if fn.endswith('.proto'):
                protos.append(os.path.join(dirpath, fn))
    protos.sort()
    return protos


def main():
    if not os.path.isdir(OUT_DIR):
        raise SystemExit('输出目录不存在: ' + OUT_DIR)
    protos = discover_protos()
    if not protos:
        raise SystemExit('未在 {0} 下发现任何 .proto'.format(PROTO_ROOT))
    for p in protos:
        build_one(p)
    print('完成：编译 {0} 个 .proto -> {1}'.format(len(protos), OUT_DIR))


if __name__ == '__main__':
    main()