import os
import re


def get_next_version(root_path: str, prefix: str) -> int:
    if not os.path.exists(root_path):
        return 0
    last = -1
    pattern = rf"\b{re.escape(prefix)}(\d+)\b"
    dirs = os.listdir(root_path)
    for dir in dirs:
        matches = re.findall(pattern, dir)
        if not matches:
            continue
        version = int(matches[0])
        last = version if version > last else last
    next = last + 1
    return next


def next_dir(root_path: str, prefix: str):
    return f"{prefix}{get_next_version(root_path, prefix=prefix)}"


def get_log_dir(name: str, root_dir: str = "logs/", prefix: str = "v") -> str:
    """Return the logdir for the model"""
    log_dir = os.path.join(root_dir, name)
    log_dir = os.path.join(log_dir, next_dir(log_dir, prefix=prefix))
    return log_dir


def get_ckpts_path(log_dir: str, ckpt_dir: str = None) -> (str, str):
    """Returns paths to last and best ckpt"""
    root = os.path.join(log_dir, ckpt_dir) if ckpt_dir is not None else log_dir
    path_best_ckpt = os.path.join(root, "best.tar")
    path_last_ckpt = os.path.join(root, "last.tar")
    return path_last_ckpt, path_best_ckpt


def get_config_path(log_dir: str, cfg_file: str = "config.yaml") -> str:
    """Returns the config path"""
    path = os.path.join(log_dir, cfg_file)
    return path
