scripts/import_snapshot.py

"""Import a snapshot into Supabase (or any Postgres + S3-compatible storage).

Reads:
    seed_data/snapshot/metadata.jsonl
    seed_data/snapshot/images/<key>

Uploads images to Supabase Storage via REST API, inserts rows with the
pre-computed embeddings into the sightings table. No model inference at all.

Run from the host with the cloud creds set:
    docker compose exec ^
        -e DATABASE_URL="postgresql+psycopg://...supabase.co:5432/postgres" ^
        -e SUPABASE_URL="https://<ref>.supabase.co" ^
        -e SUPABASE_KEY="sb_secret_..." ^
        backend python -m scripts.import_snapshot
"""
import argparse
import json
import logging
import os
from datetime import datetime
from pathlib import Path
from uuid import UUID

import requests
from sqlalchemy import create_engine
from sqlalchemy.orm import Session

from app.models import Sighting

log = logging.getLogger("import")


def upload_to_supabase_storage(
    project_url: str,
    bucket: str,
    key: str,
    body: bytes,
    auth_key: str,
    content_type: str = "image/jpeg",
    timeout: int = 30,
) -> None:
    """Upload a single object via Supabase Storage REST. Idempotent (x-upsert)."""
    url = f"{project_url.rstrip('/')}/storage/v1/object/{bucket}/{key}"
    headers = {
        "Authorization": f"Bearer {auth_key}",
        "Content-Type": content_type,
        "x-upsert": "true",
    }
    resp = requests.post(url, headers=headers, data=body, timeout=timeout)
    if resp.status_code not in (200, 201):
        raise RuntimeError(
            f"Upload {key} failed: HTTP {resp.status_code} {resp.text[:200]}"
        )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--snapshot", type=Path, default=Path("/seed_data/snapshot"))
    parser.add_argument("--bucket", default="dogs")
    parser.add_argument(
        "--skip-upload",
        action="store_true",
        help="Skip uploading images (rows only). Useful for re-running after a partial DB failure.",
    )
    parser.add_argument(
        "--skip-rows",
        action="store_true",
        help="Skip inserting DB rows (uploads only).",
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")

    project_url = os.environ.get("SUPABASE_URL")
    auth_key = os.environ.get("SUPABASE_KEY")
    db_url = os.environ.get("DATABASE_URL")

    missing = [
        name for name, val in [
            ("SUPABASE_URL", project_url),
            ("SUPABASE_KEY", auth_key),
            ("DATABASE_URL", db_url),
        ] if not val
    ]
    if missing:
        raise SystemExit(f"Missing env vars: {', '.join(missing)}")

    metadata_path = args.snapshot / "metadata.jsonl"
    images_root = args.snapshot / "images"

    if not metadata_path.exists():
        raise SystemExit(f"No metadata at {metadata_path}. Run export_snapshot first.")

    log.info("Reading %s", metadata_path)
    rows: list[dict] = []
    with metadata_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    log.info("Loaded %d rows from snapshot.", len(rows))

    # ---- Upload images ----
    if not args.skip_upload:
        total_files = sum(2 for _ in rows)  # original + cropped per row
        uploaded = 0
        for row in rows:
            for key in (row["image_key"], row["cropped_key"]):
                local_path = images_root / key
                if not local_path.exists():
                    log.warning("Missing image file: %s", local_path)
                    continue
                upload_to_supabase_storage(
                    project_url,
                    args.bucket,
                    key,
                    local_path.read_bytes(),
                    auth_key,
                )
                uploaded += 1
                if uploaded % 50 == 0:
                    log.info("Uploaded %d/%d files...", uploaded, total_files)
        log.info("Uploaded %d files to bucket '%s'.", uploaded, args.bucket)
    else:
        log.info("--skip-upload set; skipping image upload.")

    # ---- Insert DB rows ----
    if not args.skip_rows:
        log.info("Connecting to DB and inserting %d rows...", len(rows))
        engine = create_engine(db_url)
        with Session(engine) as session:
            for row in rows:
                sighting = Sighting(
                    id=UUID(row["id"]),
                    image_url=row["image_key"],
                    cropped_url=row["cropped_key"],
                    embedding=row["embedding"],
                    latitude=row["latitude"],
                    longitude=row["longitude"],
                    sighted_at=datetime.fromisoformat(row["sighted_at"]),
                    notes=row.get("notes"),
                    source=row["source"],
                    identity=row.get("identity"),
                )
                # merge() = upsert by primary key, so this script is re-runnable.
                session.merge(sighting)
            session.commit()
        log.info("Inserted %d rows.", len(rows))
    else:
        log.info("--skip-rows set; skipping DB insert.")

    log.info("Done.")


if __name__ == "__main__":
    main()