#!/usr/bin/env python3
"""
SneakerPicks Cutout Processor v3
Uses rembg HTTP service for background removal.
Smart: skips images that already have transparency.
Priority queue: P1 (top brands, high offers), P2 (medium), P3 (rest).
"""
import psycopg2
import psycopg2.extras
import requests
import os
import sys
import time
import argparse
from pathlib import Path
from io import BytesIO

try:
    from PIL import Image
except ImportError:
    print("Pillow required: pip install Pillow")
    sys.exit(1)

DB_URL = os.environ.get(
    "DATABASE_URL",
    "postgresql://supabase_admin:HzWrn2rMMxUYWO6OT5X0PScIVEvSjQMup9YSLIrkPQ@172.18.0.5:5432/sneakerpicks"
)
REMBG_URL = os.environ.get("REMBG_URL", "http://localhost:5100/api/remove")
CUTOUT_DIR = Path("/var/www/sneakerpicks/public/cutouts")
CUTOUT_URL_PREFIX = "/cutouts"

# Top brands get P1 priority
TOP_BRANDS = [
    'nike', 'adidas', 'new balance', 'jordan', 'puma', 'asics',
    'reebok', 'converse', 'vans', 'on running', 'hoka', 'salomon',
    'diadora', 'saucony', 'karhu'
]


def has_alpha(image_data: bytes) -> bool:
    """Check if image data is PNG with alpha channel."""
    try:
        img = Image.open(BytesIO(image_data))
        return img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info)
    except Exception:
        return False


def download_image(url: str, timeout: int = 30) -> bytes | None:
    """Download image, return bytes or 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 wait_for_rembg(max_wait: int = 120) -> bool:
    """Wait for rembg service to be ready."""
    import time as _time
    for i in range(max_wait // 5):
        try:
            r = requests.get(REMBG_URL.replace("/remove", "/"), timeout=5)
            return True
        except Exception:
            pass
        try:
            r = requests.post(REMBG_URL, timeout=5)
            if r.status_code in (422, 200):
                return True
        except Exception:
            pass
        if i > 0 and i % 6 == 0:
            print(f"    ⏳ Waiting for rembg service... ({i*5}s)")
        _time.sleep(5)
    return False


def remove_bg(image_data: bytes, model: str = "birefnet-general", retries: int = 3) -> bytes | None:
    """Send image to rembg service, return PNG bytes. Retries on crash."""
    for attempt in range(retries):
        try:
            r = requests.post(
                REMBG_URL,
                files={"file": ("image.jpg", image_data, "image/jpeg")},
                data={"model": model},
                timeout=180,
            )
            r.raise_for_status()
            return r.content
        except Exception as e:
            print(f"    ✗ rembg failed (attempt {attempt+1}/{retries}): {e}")
            if attempt < retries - 1:
                print(f"    ⏳ Waiting for rembg to recover...")
                time.sleep(30)
                if not wait_for_rembg():
                    print(f"    ✗ rembg service not responding")
                    return None
    return None


def get_priority_query(priority: str) -> str:
    """Build SQL query based on priority tier."""
    # Build brand filter for P1
    brand_conditions = " OR ".join(
        f"LOWER(p.brand) = '{b}'" for b in TOP_BRANDS
    )

    if priority == "p1":
        # P1: Top brands OR high offer count (5+) — most visible products
        where_extra = f"""
          AND (
            ({brand_conditions})
            OR (SELECT COUNT(*) FROM offers o WHERE o.product_id = p.id) >= 5
          )
        """
        order = "offer_count DESC, p.clicks_7d DESC NULLS LAST, p.id"
    elif priority == "p2":
        # P2: Products with some activity or medium offer count, excluding P1
        where_extra = f"""
          AND NOT (
            ({brand_conditions})
            OR (SELECT COUNT(*) FROM offers o WHERE o.product_id = p.id) >= 5
          )
          AND (
            (SELECT COUNT(*) FROM offers o WHERE o.product_id = p.id) >= 2
            OR COALESCE(p.clicks_7d, 0) > 0
          )
        """
        order = "offer_count DESC, p.clicks_7d DESC NULLS LAST, p.id"
    elif priority == "p3":
        # P3: Everything else (1 offer, no clicks, non-top brand)
        where_extra = f"""
          AND NOT (
            ({brand_conditions})
            OR (SELECT COUNT(*) FROM offers o WHERE o.product_id = p.id) >= 5
          )
          AND NOT (
            (SELECT COUNT(*) FROM offers o WHERE o.product_id = p.id) >= 2
            OR COALESCE(p.clicks_7d, 0) > 0
          )
        """
        order = "p.id"
    else:
        # "all" — process everything, but prioritize P1 first
        where_extra = ""
        order = f"""
          CASE
            WHEN ({brand_conditions})
              OR (SELECT COUNT(*) FROM offers o2 WHERE o2.product_id = p.id) >= 5
            THEN 1
            WHEN (SELECT COUNT(*) FROM offers o3 WHERE o3.product_id = p.id) >= 2
              OR COALESCE(p.clicks_7d, 0) > 0
            THEN 2
            ELSE 3
          END,
          offer_count DESC, p.clicks_7d DESC NULLS LAST, p.id
        """

    return f"""
        SELECT p.id, p.image_url, p.slug,
               (SELECT COUNT(*) FROM offers o WHERE o.product_id = p.id) as offer_count
        FROM products p
        WHERE p.cutout_image_url IS NULL
          AND p.image_url IS NOT NULL
          AND p.image_url != ''
          AND p.is_sneaker = true
          {where_extra}
        ORDER BY {order}
    """


def main():
    parser = argparse.ArgumentParser(description="Process product cutouts")
    parser.add_argument("--batch", type=int, default=50, help="Batch size")
    parser.add_argument("--model", default="isnet-general-use", help="rembg model (isnet-general-use=fast, birefnet-general=quality)")
    parser.add_argument("--delay", type=float, default=1.0, help="Delay between requests (sec)")
    parser.add_argument("--skip-transparent", action="store_true", help="Skip images that already have transparency")
    parser.add_argument("--priority", default="all", choices=["p1", "p2", "p3", "all"],
                        help="Priority tier: p1=top brands/high offers, p2=medium, p3=rest, all=everything")
    args = parser.parse_args()

    print(f"=== SneakerPicks Cutout Processor v3 ===")
    print(f"  Model: {args.model} | Batch: {args.batch} | Priority: {args.priority} | Delay: {args.delay}s")

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

    conn = psycopg2.connect(DB_URL)
    cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)

    query = get_priority_query(args.priority)
    cur.execute(query + f" LIMIT {args.batch}")

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

    processed = 0
    skipped = 0
    errors = 0
    start_time = time.time()

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

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

        # Check if cutout already exists on disk
        if cutout_path.exists() and cutout_path.stat().st_size > 0:
            cur.execute(
                "UPDATE products SET cutout_image_url = %s, needs_cutout = false WHERE id = %s",
                (cutout_url, pid),
            )
            conn.commit()
            print(f"    → Already on disk, linked")
            skipped += 1
            continue

        # Download source image
        img_data = download_image(image_url)
        if not img_data:
            cur.execute("UPDATE products SET needs_cutout = false WHERE id = %s", (pid,))
            conn.commit()
            errors += 1
            continue

        # Smart check: already transparent?
        if has_alpha(img_data):
            # Save directly as cutout
            cutout_path.write_bytes(img_data)
            cur.execute(
                "UPDATE products SET cutout_image_url = %s, needs_cutout = false WHERE id = %s",
                (cutout_url, pid),
            )
            conn.commit()
            print(f"    → Already transparent, saved directly")
            skipped += 1
            continue

        if args.skip_transparent:
            skipped += 1
            continue

        # Send to rembg
        t0 = time.time()
        result = remove_bg(img_data, args.model)
        elapsed = time.time() - t0

        if result and len(result) > 1000:  # sanity check
            cutout_path.write_bytes(result)
            cur.execute(
                "UPDATE products SET cutout_image_url = %s, needs_cutout = false WHERE id = %s",
                (cutout_url, pid),
            )
            conn.commit()
            print(f"    ✓ Done ({elapsed:.1f}s, {len(result)//1024}KB)")
            processed += 1
        else:
            print(f"    ✗ rembg failed ({elapsed:.1f}s) - will retry later")
            errors += 1

        time.sleep(args.delay)

    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()