import argparse
import csv
import shutil
import zipfile
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(
        description="Package extracted frame folders into per-video ZIP archives for CVAT."
    )
    parser.add_argument("--frames-root", required=True, help="Root with extracted frames, e.g. D:\\MAI_CVAT_frames_s15\\frames_all")
    parser.add_argument("--output-root", required=True, help="Where ZIP archives and manifests will be written")
    parser.add_argument("--group-filter", action="append", default=[], help="Optional substring filter for group names")
    parser.add_argument("--max-files-per-zip", type=int, default=0, help="Optional limit; skip items above this count when > 0")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    frames_root = Path(args.frames_root).expanduser().resolve()
    output_root = Path(args.output_root).expanduser().resolve()

    if not frames_root.exists():
        raise SystemExit(f"Frames root not found: {frames_root}")

    filters = [s.strip().lower() for s in args.group_filter if s.strip()]
    zip_root = output_root / "zips"
    zip_root.mkdir(parents=True, exist_ok=True)

    rows = []
    video_dirs = list(find_video_dirs(frames_root))
    for video_dir in video_dirs:
        rel_parts = video_dir.relative_to(frames_root).parts
        if len(rel_parts) < 5:
            continue

        split, role, group_name, video_name, interval_name = rel_parts[:5]
        if filters and not any(token in group_name.lower() for token in filters):
            continue

        files = sorted(
            [
                p for p in video_dir.iterdir()
                if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png"}
            ]
        )
        file_count = len(files)
        if file_count == 0:
            continue
        if args.max_files_per_zip > 0 and file_count > args.max_files_per_zip:
            continue

        safe_group = sanitize(group_name)
        safe_video = sanitize(video_name)
        safe_interval = sanitize(interval_name)
        zip_name = f"{safe_group}__{safe_video}__{safe_interval}.zip"
        zip_path = zip_root / zip_name

        if zip_path.exists():
            if args.overwrite:
                zip_path.unlink()
            else:
                rows.append(build_row(split, role, group_name, video_name, interval_name, file_count, zip_path, video_dir))
                continue

        with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_STORED) as zf:
            for file_path in files:
                zf.write(file_path, arcname=file_path.name)

        rows.append(build_row(split, role, group_name, video_name, interval_name, file_count, zip_path, video_dir))

    manifest_path = output_root / "cvat_task_packages.csv"
    write_csv(
        manifest_path,
        rows,
        [
            "split",
            "role",
            "group_name",
            "video_name",
            "interval_name",
            "frame_count",
            "zip_path",
            "source_dir",
        ],
    )
    print(f"Wrote {len(rows)} package rows to {manifest_path}")
    print(f"ZIP archives are in {zip_root}")


def find_video_dirs(frames_root: Path):
    for split_dir in sorted(p for p in frames_root.iterdir() if p.is_dir()):
        for role_dir in sorted(p for p in split_dir.iterdir() if p.is_dir()):
            for group_dir in sorted(p for p in role_dir.iterdir() if p.is_dir()):
                for video_dir in sorted(p for p in group_dir.iterdir() if p.is_dir()):
                    for interval_dir in sorted(p for p in video_dir.iterdir() if p.is_dir()):
                        yield interval_dir


def sanitize(value: str):
    out = []
    for ch in value:
        if ch in '<>:"/\\|?*':
            out.append("_")
        else:
            out.append(ch)
    return "".join(out).strip().rstrip(".")


def build_row(split, role, group_name, video_name, interval_name, file_count, zip_path: Path, source_dir: Path):
    return {
        "split": split,
        "role": role,
        "group_name": group_name,
        "video_name": video_name,
        "interval_name": interval_name,
        "frame_count": file_count,
        "zip_path": str(zip_path),
        "source_dir": str(source_dir),
    }


def write_csv(path: Path, rows, fieldnames):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8-sig") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


if __name__ == "__main__":
    main()