import cv2
import numpy as np
import sys
import os

def log_message(message):
    """Helper function to log messages to stderr"""
    print(message, file=sys.stderr)

def remove_background(input_path, bg_path, output_path):
    try:
        log_message(f"Memulai proses dengan file: {input_path}")
        
        # Baca gambar input
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"File input tidak ditemukan: {input_path}")
            
        img = cv2.imread(input_path)
        if img is None:
            raise ValueError(f"Gagal membaca file gambar: {input_path}")
            
        original = img.copy()
        h, w = img.shape[:2]
        log_message(f"Dimensi gambar: {w}x{h}")
        
        # Baca background
        if not os.path.exists(bg_path):
            raise FileNotFoundError(f"File background tidak ditemukan: {bg_path}")
            
        new_bg = cv2.imread(bg_path)
        if new_bg is None:
            raise ValueError(f"Gagal membaca file background: {bg_path}")
        
        new_bg = cv2.resize(new_bg, (w, h))
        
        # Konversi ke grayscale untuk masking sederhana
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Thresholding adaptif untuk masking
        mask = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, 
                                    cv2.THRESH_BINARY_INV, 21, 10)
        
        # Perbaiki mask dengan morfologi
        kernel = np.ones((5,5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        
        # Gabungkan gambar
        foreground = cv2.bitwise_and(original, original, mask=mask)
        background = cv2.bitwise_and(new_bg, new_bg, mask=cv2.bitwise_not(mask))
        result = cv2.add(foreground, background)
        
        # Simpan hasil
        cv2.imwrite(output_path, result)
        log_message(f"Proses berhasil. Hasil disimpan di: {output_path}")
        
    except Exception as e:
        log_message(f"ERROR: {str(e)}")
        raise

if __name__ == "__main__":
    try:
        if len(sys.argv) != 4:
            raise ValueError("Usage: python remove_bg.py <input_path> <bg_path> <output_path>")
        
        input_path = sys.argv[1]
        bg_path = sys.argv[2]
        output_path = sys.argv[3]
        
        # Verifikasi ekstensi output
        if not output_path.lower().endswith(('.png', '.jpg', '.jpeg')):
            raise ValueError("Format output harus PNG atau JPG")
        
        remove_background(input_path, bg_path, output_path)
        
    except Exception as e:
        print(f"ERROR: {str(e)}", file=sys.stderr)
        sys.exit(1)