import cv2
import numpy as np
import os

def slice_stream_assets():
    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. CLEAN BACKGROUND (GrabCut)
    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]
    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') 
    
    # Apply Alpha
    b, g, r = cv2.split(img)
    alpha = mask_final * 255
    img_rgba = cv2.merge([b, g, r, alpha])
    
    # 2. CREATE 'HEAD' (Pipe Only)
    # The pipe is roughly the top 40%? Let's crop visually based on the image structure.
    # The pipe is top-right.
    # Let's say we keep the top half as the "Start".
    # And we take a middle slice as the "Segment".
    
    # Finding the pipe:
    # Based on the image, the pipe drain is at the top.
    # Let's crop `head` from y=0 to y=0.45*h
    cut_y = int(h * 0.45)
    
    head_img = img_rgba[0:cut_y, :]
    
    # Crop transparent borders from head
    coords = cv2.findNonZero(head_img[:,:,3]) # Alpha
    if coords is not None:
        x, y, cw, ch = cv2.boundingRect(coords)
        head_img = head_img[y:y+ch, x:x+cw]
        
    # 3. CREATE 'SEGMENT' (Water Channel)
    # We take a slice from the middle-bottom.
    # Crop from y=0.45*h to y=0.85*h (skip very bottom tip?)
    # Actually, let's take a nice chunk that can be tiled.
    # The channel is diagonal. Tiling diagonal is hard without overlap.
    # Let's just crop the rest of the stream as one big piece for now.
    
    body_img = img_rgba[cut_y:, :]
    
    # Crop transparent borders from body
    coords = cv2.findNonZero(body_img[:,:,3])
    if coords is not None:
        x, y, cw, ch = cv2.boundingRect(coords)
        body_img = body_img[y:y+ch, x:x+cw]
        
    # 4. SOFTEN EDGES (To fix floating walls)
    # Applied to both Head and Body.
    # We want to fade out the BOTTOM edge of the mask, so the "wall" blends into the grass.
    
    def soften_bottom_edge(image):
        h, w = image.shape[:2]
        # Create a gradient mask for the bottom 20 pixels
        grad_h = 30
        if h < grad_h: return image # Too small
        
        # We modify the alpha channel
        alpha = image[:,:,3]
        
        # We need to detect where the "bottom" of the object is.
        # Since it's diagonal, it's tricky.
        # Simple hack: Erode the alpha slightly to sharpen the cut, then blur it?
        # Or just blur the edges?
        
        # Let's try blurring the alpha channel to soften the hard cut against the grass.
        # Only on the edges.
        # Get edge mask
        edges = cv2.Canny(alpha, 100, 200)
        # Dilate edges to get a rim
        rim = cv2.dilate(edges, np.ones((5,5),np.uint8))
        
        # Blur alpha
        blurred_alpha = cv2.GaussianBlur(alpha, (7,7), 0)
        
        # Apply blurred alpha where rim is
        # image[:,:,3] = np.where(rim>0, blurred_alpha, alpha)
        
        # Actually, let's just do a global soft outline.
        image[:,:,3] = blurred_alpha
        return image

    # head_img = soften_bottom_edge(head_img)
    # body_img = soften_bottom_edge(body_img)
    # (Skipping blur for now, plain cut is cleaner if geometry is right)

    # Save
    base_dir = '/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Environment'
    if not os.path.exists(base_dir): os.makedirs(base_dir, exist_ok=True)
    
    cv2.imwrite(os.path.join(base_dir, 'stream_head.png'), head_img)
    cv2.imwrite(os.path.join(base_dir, 'stream_body.png'), body_img)
    
    print("Sliced stream into head and body.")

if __name__ == "__main__":
    slice_stream_assets()