#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Local rembg processing — runs on local PC, uploads cutouts to VPS.
birefnet-general model loads once in memory, processes all images.
"""
import sys, io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace', line_buffering=True)
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace', line_buffering=True)
import psycopg2
import psycopg2.extras
import requests
import os
import sys
import time
import argparse
import subprocess
from pathlib import Path
from io import BytesIO

from PIL import Image
from rembg import remove, new_session

# DB via SSH tunnel or direct (set up SSH tunnel first: ssh -L 5433:172.18.0.5:5432 root@72.62.45.83)
DB_URL = os.environ.get(
    "DATABASE_URL",
    "postgresql://supabase_admin:HzWrn2rMMxUYWO6OT5X0PScIVEvSjQMup9YSLIrkPQ@localhost:5433/sneakerpicks"
)

VPS_HOST = "root@72.62.45.83"
VPS_KEY = os.path.expanduser("~/.ssh/hostinger_vps")
VPS_CUTOUT_DIR = "/var/www/sneakerpicks/public/cutouts"
CUTOUT_URL_PREFIX = "/cutouts"
LOCAL_CUTOUT_DIR = Path("D:/sneakerpicks-cutouts")

PRIORITY_BRANDS = {
    'nike', 'adidas', 'adidas originals', 'new balance', 'puma', 'asics',
    'jordan', 'reebok', 'converse', 'vans', 'saucony', 'hoka', 'on running',
    'salomon', 'filling pieces', 'autry', 'common projects',
}


def download_image(url: str, timeout: int = 30) -> bytes | None:
    try:
        r = requests.get(url, timeout=timeout, headers={"User-Agent": "SneakerPicks/1.0"})
        r.raise_for_status()
        return r.content
    except Exception as e:
        print(f"    ✗ Download failed: {e}")
        return None


def upload_to_vps(local_path: Path, remote_path: str) -> bool:
    """Upload file to VPS via SCP."""
    try:
        result = subprocess.run(
            ["scp", "-i", VPS_KEY, "-o", "StrictHostKeyChecking=no",
             str(local_path), f"{VPS_HOST}:{remote_path}"],
            capture_output=True, timeout=30
        )
        return result.returncode == 0
    except Exception as e:
        print(f"    ✗ Upload failed: {e}")
        return False


def batch_upload_to_vps(local_dir: Path, batch_files: list[str]) -> bool:
    """Upload batch of files via rsync for efficiency."""
    if not batch_files:
        return True
    try:
        result = subprocess.run(
            ["scp", "-i", VPS_KEY, "-o", "StrictHostKeyChecking=no"] +
            [str(local_dir / f) for f in batch_files] +
            [f"{VPS_HOST}:{VPS_CUTOUT_DIR}/"],
            capture_output=True, timeout=120
        )
        return result.returncode == 0
    except Exception as e:
        print(f"    ✗ Batch upload failed: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="Local rembg processing + VPS upload")
    parser.add_argument("--batch", type=int, default=500, help="Batch size")
    parser.add_argument("--model", default="birefnet-general", help="rembg model")
    parser.add_argument("--delay", type=float, default=0.5, help="Delay between images (sec)")
    parser.add_argument("--priority", default="all", choices=["p1", "p2", "p3", "all"])
    parser.add_argument("--upload-every", type=int, default=50, help="Upload batch every N images")
    args = parser.parse_args()

    LOCAL_CUTOUT_DIR.mkdir(parents=True, exist_ok=True)

    print(f"=== SneakerPicks Local Cutout Processor ===")
    print(f"  Model: {args.model} | Batch: {args.batch} | Delay: {args.delay}s")
    print(f"  Local dir: {LOCAL_CUTOUT_DIR}")
    print(f"  Upload every: {args.upload_every} images")
    print(f"  Loading model into memory...")

    t0 = time.time()
    session = new_session(args.model)
    print(f"  Model loaded in {time.time()-t0:.1f}s")

    print(f"  Connecting to DB via SSH tunnel...")
    conn = psycopg2.connect(DB_URL)
    cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)

    # Build priority query
    priority_filter = ""
    if args.priority == "p1":
        brands_str = ",".join(f"'{b}'" for b in PRIORITY_BRANDS)
        priority_filter = f"AND (brand_name IN ({brands_str}) OR offer_count >= 5)"
    elif args.priority == "p2":
        priority_filter = "AND offer_count >= 2"

    cur.execute(f"""
        WITH ranked AS (
            SELECT p.id, p.image_url, p.slug,
                   LOWER(COALESCE(b.name, '')) as brand_name,
                   (SELECT COUNT(*) FROM product_offers o WHERE o.product_id = p.id) as offer_count
            FROM products p
            LEFT JOIN brands b ON p.brand_id = b.id
            WHERE p.cutout_image_url IS NULL
              AND p.image_url IS NOT NULL
              AND p.image_url != ''
              AND p.is_sneaker = true
        )
        SELECT * FROM ranked
        WHERE 1=1 {priority_filter}
        ORDER BY offer_count DESC, id
        LIMIT %s
    """, (args.batch,))

    products = cur.fetchall()
    print(f"  Found {len(products)} products to process\n")

    processed = 0
    skipped = 0
    errors = 0
    start_time = time.time()
    pending_uploads = []

    for i, p in enumerate(products, 1):
        pid = p["id"]
        image_url = p["image_url"]
        slug = p["slug"] or f"product-{pid}"
        local_path = LOCAL_CUTOUT_DIR / f"{pid}.png"
        cutout_url = f"{CUTOUT_URL_PREFIX}/{pid}.png"

        print(f"  [{i}/{len(products)}] #{pid} {slug[:50]}...")

        # Already processed locally?
        if local_path.exists() and local_path.stat().st_size > 1000:
            print(f"    → Already on disk locally")
            pending_uploads.append(f"{pid}.png")
            cur.execute(
                "UPDATE products SET cutout_image_url = %s WHERE id = %s",
                (cutout_url, pid),
            )
            conn.commit()
            skipped += 1
        else:
            # Download
            img_data = download_image(image_url)
            if not img_data:
                errors += 1
                continue

            # Process with rembg
            try:
                t1 = time.time()
                input_img = Image.open(BytesIO(img_data)).convert("RGBA")
                output_img = remove(input_img, session=session)

                # Save as PNG
                buf = BytesIO()
                output_img.save(buf, format="PNG", optimize=True)
                result = buf.getvalue()
                elapsed = time.time() - t1

                if len(result) > 1000:
                    local_path.write_bytes(result)
                    pending_uploads.append(f"{pid}.png")
                    cur.execute(
                        "UPDATE products SET cutout_image_url = %s WHERE id = %s",
                        (cutout_url, pid),
                    )
                    conn.commit()
                    print(f"    ✓ Done ({elapsed:.1f}s, {len(result)//1024}KB)")
                    processed += 1
                else:
                    print(f"    ✗ Output too small ({len(result)} bytes)")
                    errors += 1
            except Exception as e:
                print(f"    ✗ Processing failed: {e}")
                errors += 1

        # Batch upload every N images
        if len(pending_uploads) >= args.upload_every:
            print(f"\n  📤 Uploading {len(pending_uploads)} cutouts to VPS...")
            if batch_upload_to_vps(LOCAL_CUTOUT_DIR, pending_uploads):
                print(f"  ✓ Upload complete")
            else:
                print(f"  ✗ Upload failed — will retry next batch")
            pending_uploads = []

        time.sleep(args.delay)

    # Upload remaining
    if pending_uploads:
        print(f"\n  📤 Final upload: {len(pending_uploads)} cutouts...")
        batch_upload_to_vps(LOCAL_CUTOUT_DIR, pending_uploads)

    total_time = time.time() - start_time
    print(f"\n{'='*50}")
    print(f"  Processed: {processed} | Skipped: {skipped} | Errors: {errors}")
    print(f"  Total time: {total_time:.0f}s | Avg: {total_time/max(processed,1):.1f}s/image")

    cur.close()
    conn.close()


if __name__ == "__main__":
    main()