import os
from pathlib import Path

#### 项目结构 ####
PROJECT_STRUCTURE = {
    # 模块文件
    "models": [
        "__init__.py",
    ],
    "training": [
        "__init__.py",
    ],
    "visualization": [
        "__init__.py",
    ],
    "utils": [
        "__init__.py",
    ],
    "configs": [
        "config.yaml",
    ],
    "scripts": [
        "run_pipeline_1.py",
        "run_pipeline_2.py",
        "run_pipeline_3.py",
    ],

    # 非模块文件
    "data": [
        "README.md",
    ],
    "logs": [],
    "outputs": [],
}

ROOT_FILES = [
    ".gitignore",
    "README.md",
    "requirements.txt",
]

#### 具体内容 ####
GITIGNORE_CONTENT = \
"""
# Python
__pycache__/
*.pyc
*.pyo
*.pyd

# virtual env
.venv/
venv/
env/

# logs
logs/

# outputs
outputs/

# dataset
data/*

# IDE
.vscode/
.idea/

# mac
.DS_Store
"""

def _make(file_path, content=None, overwrite=False):
    # 删去旧的, 创建新的
    if os.path.exists(file_path) and overwrite:
        print(f"Overwriting {file_path}...")
        file_path.unlink()
        if content:
            file_path.write_text(content.strip() + "\n")
        else:
            file_path.touch(exist_ok=True)
    # 保持不变
    elif os.path.exists(file_path):
        pass
    # 创建新的
    else:
        if content:
            file_path.write_text(content.strip() + "\n")
        else:
            file_path.touch(exist_ok=True)


def b_create_project(project_dir: Path, overwrite=False):
    """
    创建项目目录结构
    """
    project_dir = Path(project_dir)

    # 创建根目录内文件夹
    for folder, files in PROJECT_STRUCTURE.items():
        folder_path = project_dir / folder
        folder_path.mkdir(parents=True, exist_ok=True)

        for file in files:
            file_path = (folder_path / file)
            _make(file_path, overwrite=overwrite)

    # 创建根目录内文件
    for file in ROOT_FILES:
        file_path = project_dir / file

        if file == ".gitignore":
            _make(file_path, content=GITIGNORE_CONTENT, overwrite=overwrite)
        else:
            _make(file_path, overwrite=overwrite)


def main():
    base_path = Path(".").resolve()
    print(f"Creating project structure in: {base_path}")

    b_create_project(base_path)
    print("Project structure created successfully!")


if __name__ == "__main__":
    main()