import cv2
import numpy as np
import os

def clean_pipe_stream_final():
    src_path = '/Users/davidkotnik/.gemini/antigravity/brain/07019d04-a214-43ab-9565-86f4e8f17e5b/uploaded_media_1769607894587.jpg'
    
    print(f"Loading {src_path}")
    img = cv2.imread(src_path)
    if img is None: return

    # 1. GRABCUT for Clean Foreground (Pipe + Stream + Earth)
    mask = np.zeros(img.shape[:2], np.uint8)
    bgdModel = np.zeros((1,65),np.float64)
    fgdModel = np.zeros((1,65),np.float64)
    h, w = img.shape[:2]
    # Rect: slight margin
    cv2.grabCut(img, mask, (10, 10, w-20, h-20), bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
    mask_final = np.where((mask==2)|(mask==0),0,1).astype('uint8') # 0 = BG, 1 = FG
    
    # 2. FILL THE DRAIN HOLE (Make it opaque)
    # The drain is typically at the top-right or top area.
    # We can detect the "grating" (dark criss-cross lines).
    # Simple heuristic: The pipe opening is a dark area.
    # To prevent it from becoming transparent, we force the mask to be 1 in that region?
    # GrabCut usually keeps the hole opaque if it's distinct from BG.
    # But if there are "holes" in the GrabCut mask (transparency inside the object), we should fill them.
    
    # Close small holes in the mask (Morphological Closing)
    kernel = np.ones((5,5),np.uint8)
    mask_closed = cv2.morphologyEx(mask_final, cv2.MORPH_CLOSE, kernel)
    
    # 3. "BURYING" (Removing outer walls)
    # The outer walls are usually at the bottom-left and bottom-right edges of the mask.
    # We can try to "shave" the bottom of the mask to reduce the "block height".
    # Let's shift the mask up? No.
    # Let's erode the mask from the bottom?
    # We can assume the "floor" is higher.
    # This is risky, but let's try a mild erosion just on the edges to smooth the blend.
    # Actually, if we just ensure the alpha channel is clean, it might be enough.
    
    # Convert to RGBA
    b, g, r = cv2.split(img)
    alpha = mask_closed * 255
    
    # 4. COLOR FIX: If the drain hole inside is gray/white (from original image),
    # the user might want it BLACK (so it looks deep).
    # Let's darken the "dark" areas inside the foreground.
    # Convert to HSV, find dark areas in FG, make them blacker.
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    # Dark areas: Value < 50?
    dark_mask = (hsv[:,:,2] < 80) & (mask_closed == 1)
    img[dark_mask] = [20, 20, 20] # Very dark gray/black
    
    img_rgba = cv2.merge([img[:,:,0], img[:,:,1], img[:,:,2], alpha])

    # Crop
    coords = cv2.findNonZero(alpha)
    if coords is not None:
        x, y, cw, ch = cv2.boundingRect(coords)
        img_rgba = img_rgba[y:y+ch, x:x+cw]

    # Save
    targets = [
        '/Users/davidkotnik/repos/novafarma/main/assets/stream_pipe.png',
        '/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Environment/stream_pipe.png'
    ]
    for t in targets:
        cv2.imwrite(t, img_rgba)
        print(f"Saved {t}")

if __name__ == "__main__":
    clean_pipe_stream_final()