#!/usr/bin/env python3
"""
IncomeBrowser — Background Removal Script
Called by remove_bg.php
Usage: python3 rembg_process.py input.jpg output.png

Install once via SSH:
  pip3 install rembg pillow --user
  OR
  pip3 install rembg[gpu] pillow --user  (if server has GPU)
"""

import sys
import os

def remove_background(input_path: str, output_path: str) -> bool:
    try:
        from rembg import remove
        from PIL import Image
        import io

        # Open input image
        with open(input_path, 'rb') as f:
            input_data = f.read()

        # Remove background using U2Net model
        # First run downloads the model (~170MB) — subsequent runs are instant
        output_data = remove(
            input_data,
            alpha_matting=True,              # Better edge quality
            alpha_matting_foreground_threshold=240,
            alpha_matting_background_threshold=10,
            alpha_matting_erode_size=10,
        )

        # Save result
        with open(output_path, 'wb') as f:
            f.write(output_data)

        # Verify output
        img = Image.open(output_path)
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
            img.save(output_path, 'PNG')

        return True

    except ImportError:
        # rembg not installed — try onnxruntime fallback
        print("ERROR: rembg not installed. Run: pip3 install rembg pillow --user", file=sys.stderr)
        sys.exit(2)

    except Exception as e:
        print(f"ERROR: {str(e)}", file=sys.stderr)
        sys.exit(1)


def remove_background_simple(input_path: str, output_path: str) -> bool:
    """
    Fallback method using basic color detection if rembg fails.
    Not as good but works without ML models.
    """
    try:
        from PIL import Image
        import numpy as np

        img = Image.open(input_path).convert('RGBA')
        data = np.array(img)

        # Sample background color from corners
        corners = [
            data[0, 0, :3],
            data[0, -1, :3],
            data[-1, 0, :3],
            data[-1, -1, :3],
        ]
        bg_color = np.mean(corners, axis=0)

        # Create alpha mask based on color distance from background
        r, g, b = data[:,:,0], data[:,:,1], data[:,:,2]
        br, bg, bb = bg_color[0], bg_color[1], bg_color[2]

        dist = np.sqrt(
            0.3 * (r.astype(float) - br)**2 +
            0.59 * (g.astype(float) - bg)**2 +
            0.11 * (b.astype(float) - bb)**2
        )

        # Threshold and smooth
        threshold = 60
        alpha = np.clip((dist / threshold) * 255, 0, 255).astype(np.uint8)
        data[:,:,3] = alpha

        result = Image.fromarray(data, 'RGBA')
        result.save(output_path, 'PNG')
        return True

    except Exception as e:
        print(f"Fallback ERROR: {str(e)}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} input_image output_image.png", file=sys.stderr)
        sys.exit(1)

    input_path  = sys.argv[1]
    output_path = sys.argv[2]

    if not os.path.exists(input_path):
        print(f"ERROR: Input file not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    # Try AI removal first, fall back to color-based
    try:
        from rembg import remove
        success = remove_background(input_path, output_path)
    except ImportError:
        print("rembg not found, using basic removal...", file=sys.stderr)
        success = remove_background_simple(input_path, output_path)

    if success:
        sys.exit(0)
    else:
        sys.exit(1)