from dotenv import load_dotenv
load_dotenv()

import os
import subprocess
import json

data_dir = os.path.join(os.path.dirname(__file__), "data")
hf_bin = os.path.join(os.path.dirname(__file__), ".venv", "bin", "hf")
bucket_url = "hf://buckets/malcolmstone/language"

if not os.path.exists(data_dir):
    print(f"✗ No data directory found at {data_dir}")
    print("  Run main.py first to download and save datasets.")
    exit(1)

# First, do a dry run to see what needs syncing
print(f"Checking what needs syncing to {bucket_url} ...")
dry_run = subprocess.run(
    [hf_bin, "sync", data_dir, bucket_url, "--dry-run"],
    capture_output=True, text=True
)

if dry_run.returncode != 0:
    print(f"  ✗ Dry run failed: {dry_run.stderr}")
    exit(1)

# Parse dry run output to count uploads needed
uploads = 0
total_size = 0
skips = 0
for line in dry_run.stdout.strip().split("\n"):
    if not line.strip():
        continue
    try:
        entry = json.loads(line)
        if entry.get("type") == "operation" and entry.get("action") == "upload":
            uploads += 1
            total_size += entry.get("size", 0)
        elif entry.get("type") == "header":
            skips = entry.get("summary", {}).get("skips", 0)
    except json.JSONDecodeError:
        pass

if uploads == 0:
    print(f"  ✓ Already in sync! Nothing to upload. ({skips} files already up to date)")
else:
    size_mb = total_size / (1024 * 1024)
    print(f"  {uploads} files to upload ({size_mb:.1f} MB)")
    print(f"  Syncing...")
    sync_result = subprocess.run(
        [hf_bin, "sync", data_dir, bucket_url, "--verbose"]
    )
    if sync_result.returncode == 0:
        print("  ✓ Sync complete!")
    else:
        print(f"  ✗ Sync failed (exit code {sync_result.returncode})")
        exit(1)