from pathlib import Path
from datetime import datetime
import textwrap

#### 项目结构 ####
PROJECT_STRUCTURE = {
    # 模块文件
    "models": [
        "__init__.py",
    ],
    "train": [
        "__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": [],
    "logs": [],
    "outputs": [],
}

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

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

    # virtual env
    .venv/
    venv/
    env/

    # logs
    logs/

    # outputs
    outputs/

    # dataset
    data/*

    # IDE
    .vscode/
    .idea/

    # mac
    .DS_Store
""")

INIT_CONTENT = textwrap.dedent("""


    __all__ = [

    ]
""")

CONFIG_YAML_CONTENT = textwrap.dedent("""
    # =========================
    # seed
    # =========================
    seed: 2026

    # =========================
    # Paths
    # =========================
    paths:
      output_dir: outputs
      log_dir: logs

    # =========================
    # Train
    # =========================
    train:
      batch_size: 64
      epochs: 50
      learning_rate: 0.001
      weight_decay: 0.0001
""")

PIPELINE_CONTENT = textwrap.dedent("""
    def main():
        pass

    if __name__ == "__main__":
        main()
""")


def _strip_once(s: str) -> str:
    if s.startswith("\n"):
        s = s[1:]
    if s.endswith("\n"):
        s = s[:-1]
    return s


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


def _makes(file: str, file_path: Path, overwrite=False):
    if file == "__init__.py":
        _make(file_path, content=INIT_CONTENT, overwrite=overwrite)

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

    elif file == "config.yaml":
        _make(file_path, content=CONFIG_YAML_CONTENT, overwrite=overwrite)

    elif file in ["run_pipeline_1.py", "run_pipeline_2.py", "run_pipeline_3.py"]:
        _make(file_path, content=PIPELINE_CONTENT, overwrite=overwrite)

    else:
        _make(file_path, overwrite=overwrite)


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)
            _makes(file, file_path, overwrite=overwrite)

    # 创建根目录内文件
    for file in ROOT_FILES:
        file_path = project_dir / file
        _makes(file, file_path, overwrite=overwrite)

def _get_time():
    current_time = datetime.now()
    time_str = current_time.strftime("%Y_%m_%d-%H_%M_%S")
    return time_str
class _manager:
    def __init__(self, dir: Path):
        self.dir = Path(dir)

    def _get_count(self):
        # 统计self.run_dir_path下的文件夹数
        dirs = [1 for item in self.dir.iterdir() if item.is_dir()]
        count_1 = sum(dirs)

        # 统计文件夹的最大编号
        count_2 = 0
        for item in self.dir.iterdir():
            if item.is_dir():
                prefix = item.name.split("-", 1)[0]
                if prefix.isdigit():
                    count_2 = max(count_2, int(prefix))

        count = max(count_1, count_2)

        return count

    def create_dir(self, tag:str=None) -> tuple[Path, str]:
        '''
        :return: dir_path, dirname
        '''
        index = self._get_count() + 1
        time = _get_time()

        dir_path = self.dir / f'{index:03d}-{time}'
        if tag:
            dir_path = self.dir / f'{index:03d}-{time}-{tag}'
        dir_path.mkdir(parents=True, exist_ok=True)

        return dir_path, dir_path.name
class B_Project_Manager:
    def __init__(self, project_dir: Path):
        self.project_dir = Path(project_dir)

    def init_project(self, overwrite=False):
        b_create_project(self.project_dir, overwrite=overwrite)

    def get_logs(self):
        path = self.project_dir / "logs"
        assert path.exists(), "logs folder not exist"

        logs_manager = _manager(path)

        return path, logs_manager

    def get_outputs(self):
        path = self.project_dir / "outputs"
        assert path.exists(), "outputs folder not exist"

        outputs_manager = _manager(path)

        return path, outputs_manager

def main():
    # b_create_project("./")

    print(_get_time())

if __name__ == "__main__":
    main()