import cv2
import numpy as np
import os
import argparse
import sys

def clean_and_save(input_path, output_path):
    # Ensure directory exists
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    print(f"Processing: {input_path}")
    
    # Read image
    img = cv2.imread(input_path)
    if img is None:
        print(f"Error: Could not read {input_path}")
        sys.exit(1)

    # Convert to RGBA
    img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    # Convert to HSV for better color segmentation
    hsv = cv2.cvtColor(img[:,:,:3], cv2.COLOR_BGR2HSV)

    # Define Magenta range in HSV
    # Magenta is around 300 degrees. OpenCV Hue is 0-179 (degrees/2). So 300/2 = 150.
    # Range: 140 - 160 approx.
    # We want to catch the solid pink background.
    
    # Lower and Upper bounds for Magenta/Pink
    # Hue: 140-170, Saturation: > 50, Value: > 50
    lower_magenta = np.array([140, 50, 50])
    upper_magenta = np.array([170, 255, 255])
    
    mask = cv2.inRange(hsv, lower_magenta, upper_magenta)
    
    # Invert mask (keep non-magenta)
    mask_inv = cv2.bitwise_not(mask)
    
    # Update alpha channel
    # Where mask is 255 (magenta), alpha should be 0.
    # We can just set the alpha channel to mask_inv directly.
    # However, to avoid hard edges, we might want to erode/dilate, but for pixel/vector props hard cut is usually fine.
    
    # Let's perform a small morphological operation to remove pink noise on edges if any
    kernel = np.ones((1,1), np.uint8)
    mask_inv = cv2.morphologyEx(mask_inv, cv2.MORPH_OPEN, kernel)
    
    b, g, r, a = cv2.split(img)
    img = cv2.merge((b, g, r, mask_inv))
    
    # Crop to content (optional but good for assets)
    # Find bounding box of non-transparent area
    coords = cv2.findNonZero(mask_inv)
    if coords is not None:
        x, y, w, h = cv2.boundingRect(coords)
        # Add 1px padding if possible
        x = max(0, x-1)
        y = max(0, y-1)
        w = min(img.shape[1]-x, w+2)
        h = min(img.shape[0]-y, h+2)
        img = img[y:y+h, x:x+w]
    
    # Save
    cv2.imwrite(output_path, img)
    print(f"Successfully saved cleaned image to {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Remove magenta background and save.')
    parser.add_argument('input', help='Input file path')
    parser.add_argument('output', help='Output file path')
    args = parser.parse_args()

    clean_and_save(args.input, args.output)