import cv2
import numpy as np
import sys
import os

def clean_pink_and_soften(input_path, output_path):
    print(f"Processing {input_path}...")
    
    img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
    if img is None:
        print("Failed to load image.")
        return

    # Check channels
    if img.shape[2] == 3:
        b, g, r = cv2.split(img)
        # Create alpha (255)
        a = np.ones_like(b) * 255
        img = cv2.merge((b, g, r, a))

    # Targeted Pink Range in HSV
    hsv = cv2.cvtColor(img[:,:,0:3], cv2.COLOR_BGR2HSV)
    lower_pink = np.array([145, 50, 50])
    upper_pink = np.array([170, 255, 255])
    
    pink_mask = cv2.inRange(hsv, lower_pink, upper_pink)
    
    # 1. REMOVE PINK (Set alpha to 0)
    # Get current alpha
    b, g, r, a = cv2.split(img)
    
    # Where pink_mask is 255, set alpha to 0
    # Also combine with existing alpha (if image was already transparent, keep it)
    new_alpha = cv2.bitwise_and(a, cv2.bitwise_not(pink_mask))

    # 2. EDGE SOFTENING (Alpha Blur 1px)
    # We blur the alpha channel slightly to soften jagged edges
    soft_alpha = cv2.GaussianBlur(new_alpha, (3, 3), 0)
    
    # Thresholding back to keep semi-transparency only on very edge?
    # Or just keep the blur. For "Vector Style", usually we want sharp, but 
    # user asked for "zmehčaj robove (Alpha Blur: 1px)".
    # Gaussian (3,3) is roughly 1px radius.

    # 3. COLOR BLEED REMOVAL (Halo effect)
    # Often pink pixels leave a pink fringe. 
    # We can dilate the non-pink color into the transparent area?
    # Or simpler: desaturate or darken pixels on the edge of transparency.
    # For now, let's stick to the requested Blur.

    final_img = cv2.merge((b, g, r, soft_alpha))
    
    cv2.imwrite(output_path, final_img)
    print(f"Saved cleaned image to {output_path}")

if __name__ == "__main__":
    base_dir = 'assets/DEMO_FAZA1'
    target_subdirs = ['Vegetation', 'Ground', 'Environment']

    for subdir in target_subdirs:
        target_dir = os.path.join(base_dir, subdir)
        if not os.path.exists(target_dir):
            print(f"Directory not found: {target_dir}")
            continue

        print(f"--- Processing {subdir} ---")
        for filename in os.listdir(target_dir):
            if filename.lower().endswith(".png"):
                # Filter for specific files if needed, or process all to be safe as per user request "na vseh teh spritih"
                # User specifically mentioned: grass, stream_water, tla_trava_tekstura
                full_path = os.path.join(target_dir, filename)
                clean_pink_and_soften(full_path, full_path)