import os
import random
import shutil
from pathlib import Path


def split_files(
    src_dir: str,
    out_dir: str,
    ratio=(70, 30),
    seed: int = 42,
    move: bool = False,
    exts=None,  # e.g. {".jpg", ".png"} or None for all files
):
    src = Path(src_dir)
    out = Path(out_dir)
    out_a = out / "train"
    out_b = out / "vaild"
    out_a.mkdir(parents=True, exist_ok=True)
    out_b.mkdir(parents=True, exist_ok=True)

    files = []
    for p in src.iterdir():
        if p.is_file():
            if exts is None or p.suffix.lower() in exts:
                files.append(p)

    if not files:
        raise ValueError(f"No files found in {src} with exts={exts}")

    random.seed(seed)
    random.shuffle(files)

    total = len(files)
    a_count = int(total * (ratio[0] / sum(ratio)))

    a_files = files[:a_count]
    b_files = files[a_count:]

    op = shutil.move if move else shutil.copy2

    for p in a_files:
        op(str(p), str(out_a / p.name))

    for p in b_files:
        op(str(p), str(out_b / p.name))

    return len(a_files), len(b_files)


a, b = split_files(
    src_dir="./images",
    out_dir="./",
    # ratio=(70, 30),
    seed=123,
    move=False,  # set True to move instead of copy
    exts=None,  # or {".jpg", ".png"}
)

c, d = split_files(
    src_dir="./labels",
    out_dir="./",
    # ratio=(70, 30),
    seed=123,
    move=False,  # set True to move instead of copy
    exts=None,  # or {".jpg", ".png"}
)