#!/usr/bin/env python3
"""Sync Label Studio data directory and logs with R2."""

from __future__ import annotations

import argparse
import sys
from datetime import datetime, timezone
from pathlib import Path

from r2_storage import PREFIX_LS, download_prefix, sync_directory


def log(message: str, *, file: object = sys.stdout) -> None:
    timestamp = datetime.now(timezone.utc).isoformat(timespec="seconds").replace(
        "+00:00", "Z"
    )
    print(f"[{timestamp}] [r2-sync] {message}", file=file, flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Sync Label Studio state with R2")
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path("/label-studio/data"),
    )
    parser.add_argument("direction", choices=["pull", "push", "both"])
    args = parser.parse_args()
    data_dir: Path = args.data_dir

    if args.direction in ("pull", "both"):
        n = download_prefix(PREFIX_LS, data_dir)
        log(f"Pulled {n} objects from r2://…/{PREFIX_LS}/ -> {data_dir}")

    if args.direction in ("push", "both"):
        if not data_dir.exists():
            log(f"No data dir at {data_dir}", file=sys.stderr)
            sys.exit(1)
        n = sync_directory(data_dir, PREFIX_LS)
        logs = data_dir.parent / "logs"
        if logs.is_dir():
            n += sync_directory(logs, f"{PREFIX_LS}/logs")
        log(f"Pushed {n} files to r2://…/{PREFIX_LS}/")


if __name__ == "__main__":
    main()