# generate_pyi.py

import re
import os
import argparse
from pathlib import Path

def sanitize_name(name: str) -> str:
    name = re.sub(r'[^a-zA-Z0-9_]', '_', name)
    return re.sub(r'^[^a-zA-Z_]+', '', name)

def pascal_case(s: str) -> str:
    """将字符串转为 PascalCase"""
    return ''.join(word.capitalize() for word in re.split(r'[\W_]+', s) if word)

def generate_pyi(base_dir: str, output_file: str):
    base_path = Path(base_dir)
    if not base_path.exists():
        raise FileNotFoundError(f"资源目录不存在: {base_dir}")

    # 构建目录树结构
    tree = {}

    for root, dirs, files in os.walk(str(base_path)):
        relative_path = Path(root).relative_to(base_path)
        parts = list(relative_path.parts) if relative_path != Path('.') else []

        # 插入当前路径到树中
        current = tree
        for part in parts:
            safe_part = sanitize_name(part)
            if safe_part not in current:
                current[safe_part] = {}
            current = current[safe_part]

        # 添加图像文件作为叶子节点
        for file in files:
            if file.startswith('.'):
                continue
            if file.split('.')[-1].lower() in {'png', 'jpg', 'jpeg', 'bmp', 'gif'}:
                stem = sanitize_name(Path(file).stem)
                current[stem] = None  # 表示这是一个图像资源（str 类型）

    # 构建类定义
    class_defs = []
    written_classes = set()

    def build_class(class_name: str, data: dict):
        if not data or class_name in written_classes:
            return
        written_classes.add(class_name)

        lines = [f"class {class_name}:\n"]
        has_content = False

        for key in sorted(data.keys()):
            value = data[key]
            safe_key = key
            if value is None:
                # 图像文件：直接定义为 str 属性
                lines.append(f"    {safe_key}: str\n")
                has_content = True
            else:
                # 子目录：引用一个类
                nested_class_name = pascal_case(safe_key)
                lines.append(f"    {safe_key}: \"{nested_class_name}\"\n")
                build_class(nested_class_name, value)
                has_content = True

        if has_content:
            class_defs.append("".join(lines))
        else:
            # 空类不写入
            written_classes.remove(class_name)

    # 构建主类 ImageResourceLoader
    loader_lines = []
    for key in sorted(tree.keys()):
        value = tree[key]
        if not value:  # 跳过空的子目录
            loader_lines.append(f"    {key}: str\n")
            continue
        class_name = pascal_case(key)
        loader_lines.append(f"    {key}: {class_name}\n")
        build_class(class_name, value)

    # 写入 .pyi 文件
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("from typing import Optional, Any\n\n")

        # 写入所有子类定义
        for line in class_defs:
            f.write(line + "\n")

        # 写入主类
        f.write("class ImageResourceLoader:\n")
        f.write("    def __init__(self, base_dir: str): ...\n")
        f.write("    def list_resources(self) -> None: ...\n")
        f.write("    def get(self, path: str, default: Optional[str] = ...) -> Optional[str]: ...\n")
        f.write("    def __getattr__(self, name: str) -> Any: ...\n")
        f.write("\n    # 自动生成的资源类属性（仅用于 IDE 提示）\n")
        f.writelines(loader_lines)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="为 ImageResourceLoader 类生成 .pyi 存根文件。")
    parser.add_argument("source_file", help="Python 源文件路径（例如：image_loader.py）")
    parser.add_argument("--base_dir", default=".automation",
                        help="图像资源目录，默认值：.automation")
    args = parser.parse_args()

    source_path = Path(args.source_file)
    if not source_path.exists():
        raise FileNotFoundError(f"源文件不存在: {args.source_file}")

    # 输出 .pyi 文件路径
    pyi_path = source_path.with_suffix(".pyi")

    print(f"生成存根文件: {pyi_path}")
    print(f"使用资源目录: {args.base_dir}")

    generate_pyi(args.base_dir, str(pyi_path))
    """ 
    python scripts/generate_pyi.py image_loader.py --base_dir assets/images

    """