#!/usr/bin/env python3
"""
Direct rembg processing — bypasses HTTP service, loads model once in memory.
Much more stable for birefnet-general on CPU.
"""
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
    from rembg import remove, new_session
except ImportError:
    print("Required: pip install Pillow rembg")
    sys.exit(1)

DB_URL = os.environ.get(
    "DATABASE_URL",
    "postgresql://supabase_admin:HzWrn2rMMxUYWO6OT5X0PScIVEvSjQMup9YSLIrkPQ@172.18.0.5:5432/sneakerpicks"
)
CUTOUT_DIR = Path("/var/www/sneakerpicks/public/cutouts")
CUTOUT_URL_PREFIX = "/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 has_alpha(image_data: bytes) -> bool:
    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 main():
    parser = argparse.ArgumentParser(description="Direct rembg processing (no HTTP)")
    parser.add_argument("--batch", type=int, default=100, help="Batch size")
    parser.add_argument("--model", default="birefnet-general", help="rembg model")
    parser.add_argument("--delay", type=float, default=1.0, help="Delay between images (sec)")
    parser.add_argument("--priority", default="all", choices=["p1", "p2", "p3", "all"])
    args = parser.parse_args()

    print(f"=== SneakerPicks Direct Cutout Processor ===")
    print(f"  Model: {args.model} | Batch: {args.batch} | Delay: {args.delay}s")
    print(f"  Loading model into memory...")

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

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

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

    # Build priority query (use CTE so offer_count alias is available in WHERE)
    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,
                   COUNT(po.id) as offer_count,
                   COALESCE(p.clicks_7d, 0) as clicks_7d
            FROM products p
            LEFT JOIN brands b ON p.brand_id = b.id
            LEFT JOIN product_offers po ON po.product_id = p.id
            WHERE p.cutout_image_url IS NULL
              AND p.image_url IS NOT NULL
              AND p.image_url != ''
              AND p.is_sneaker = true
            GROUP BY p.id, b.name
        )
        SELECT * FROM ranked
        WHERE 1=1 {priority_filter}
        ORDER BY clicks_7d DESC, 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()

    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]}...")

        # Already 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
        img_data = download_image(image_url)
        if not img_data:
            errors += 1
            continue

        # Already transparent?
        if has_alpha(img_data):
            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

        # Process with rembg directly in-memory
        try:
            t1 = time.time()
            input_img = Image.open(BytesIO(img_data))
            output_img = remove(input_img, session=session)
            
            # Save as PNG
            buf = BytesIO()
            output_img.save(buf, format="PNG")
            result = buf.getvalue()
            elapsed = time.time() - t1

            if len(result) > 1000:
                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"    ✗ Output too small ({len(result)} bytes)")
                errors += 1
        except Exception as e:
            print(f"    ✗ Processing failed: {e}")
            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()