#!/usr/bin/env python3
#
# @File: upload-training-images.py
# @Date: 2026-06-02
#
# Uploads web-scraped training images to the media server so ImgIX can
# process them via fetch-imgix-masks.py.
#
# Web images have no server path — this script uploads them to a dedicated
# training directory and writes the paths.txt that fetch-imgix-masks.py needs.
#
# Usage:
#   python3 tools/internal/upload-training-images.py
#   python3 tools/internal/upload-training-images.py --source candidates/filtered

import argparse
import subprocess
import sys
from pathlib import Path

SSH_SERVER  = "root@165.227.32.132"
MEDIA_PATH  = "/var/www/media.liftkit.click/public_html"
REMOTE_DIR  = "training-candidates"          # subdir on server
SSH_KEY     = str(Path.home() / ".ssh" / "id_ed25519")
SSH_OPTS    = [
    "-i", SSH_KEY,
    "-o", "StrictHostKeyChecking=no",
    "-o", "ConnectTimeout=15",
]


def ssh(cmd: str) -> str:
    r = subprocess.run(
        ["ssh"] + SSH_OPTS + [SSH_SERVER, cmd],
        capture_output=True, text=True, timeout=30,
    )
    if r.returncode != 0:
        print(f"SSH error: {r.stderr.strip()}")
        sys.exit(1)
    return r.stdout.strip()


def scp_up(local: Path, remote: str) -> bool:
    r = subprocess.run(
        ["scp"] + SSH_OPTS + [str(local), f"{SSH_SERVER}:{remote}"],
        capture_output=True, text=True, timeout=60,
    )
    return r.returncode == 0


def main():
    parser = argparse.ArgumentParser(description="Upload filtered images to media server for ImgIX")
    parser.add_argument("--source", default="candidates/filtered", help="Folder of filtered JPEGs")
    args = parser.parse_args()

    source_dir = Path(args.source)
    images     = sorted(source_dir.glob("*.jpg"))
    if not images:
        print(f"No JPEGs in {source_dir}")
        sys.exit(1)

    remote_base = f"{MEDIA_PATH}/{REMOTE_DIR}"
    paths_file  = source_dir / "paths.txt"

    # Load existing to support resume
    path_map = {}
    if paths_file.exists():
        for line in paths_file.read_text().splitlines():
            if "\t" in line:
                stem, rel = line.split("\t", 1)
                path_map[stem] = rel.strip()

    already_done = set(path_map.keys())
    remaining    = [p for p in images if p.stem not in already_done]

    print(f"Images to upload : {len(remaining)}  ({len(already_done)} already done)")
    if not remaining:
        print("All images already uploaded. paths.txt is up to date.")
        print(f"Next step:")
        print(f"  python3 tools/internal/fetch-imgix-masks.py --source {source_dir}")
        return

    # Ensure remote directory exists
    ssh(f"mkdir -p {remote_base}")
    print(f"Remote dir : {SSH_SERVER}:{remote_base}\n")

    ok = fail = 0
    for img_path in remaining:
        remote_path = f"{remote_base}/{img_path.name}"
        rel_path    = f"{REMOTE_DIR}/{img_path.name}"
        print(f"  {img_path.name}", end=" ... ", flush=True)
        if scp_up(img_path, remote_path):
            path_map[img_path.stem] = rel_path
            print(f"OK")
            ok += 1
        else:
            print("FAILED")
            fail += 1

    # Write paths.txt
    with open(paths_file, "w") as f:
        for stem, rel in sorted(path_map.items()):
            f.write(f"{stem}\t{rel}\n")

    print(f"\n{'='*50}")
    print(f"Uploaded  : {ok}")
    print(f"Failed    : {fail}")
    print(f"paths.txt : {paths_file}")
    print(f"\nNext step:")
    print(f"  python3 tools/internal/fetch-imgix-masks.py --source {source_dir}")
    print("=" * 50)


if __name__ == "__main__":
    main()