import cv2
import numpy as np
import os

# Paths
input_path = '/Users/davidkotnik/.gemini/antigravity/brain/8233d64e-0c17-43b1-b8b5-fbc41754e56b/uploaded_media_1769679216460.jpg'
output_dir = '/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Environment/'

# Ensure output dir exists
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# Read image
img = cv2.imread(input_path, cv2.IMREAD_COLOR)
if img is None:
    print("Error loading input image")
    exit(1)

h, w = img.shape[:2]
print(f"Image loaded: {w}x{h}")

# Convert to BGRA
img_rgba = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

# --- 1. REMOVE BACKGROUND (Checkerboard/White) ---
# Strategy: The background is bright (White/Light Gray). The content (Mud, Pipe, Debris) is darker.
# Threshold on Value (V) channel.
# Let's verify histogram or min/max.
# Or just a hard threshold.
# Mud is dark. Pipe is dark gray.
# Debris is dark.
# Checkerboard is > 200 usually.

thresh_val = 190
# Create mask where Value > 190 -> Background
_, mask_bg_bright = cv2.threshold(hsv[:, :, 2], thresh_val, 255, cv2.THRESH_BINARY)

# Also check for Low Saturation for the checkerboard?
# Checkerboard is grayscale (S near 0).
# Pipe is also grayscale.
# But Pipe is Darker.
# So Brightness threshold should handle it.

# Refine mask:
# Remove small holes in background mask?
# Use morphological open
kernel = np.ones((3,3), np.uint8)
mask_bg_bright = cv2.morphologyEx(mask_bg_bright, cv2.MORPH_OPEN, kernel)

# --- 2. REMOVE GREEN TRIANGLE ---
# Green range
lower_green = np.array([35, 40, 40])
upper_green = np.array([90, 255, 255])
mask_green = cv2.inRange(hsv, lower_green, upper_green)

# Combine Background and Green masks
# We want to remove (set alpha 0) where mask is BG or Green.
mask_remove = cv2.bitwise_or(mask_bg_bright, mask_green)

# Set alpha
# Start with opaque (255)
img_rgba[:, :, 3] = 255
# Set 0 where removal mask is active
img_rgba[mask_remove > 0, 3] = 0

# Smooth edges slightly
# cv2.GaussianBlur on alpha channel?
# Or erode the alpha slightly to remove fringes.
# Let's erode the opaque region.
alpha = img_rgba[:, :, 3]
alpha = cv2.erode(alpha, kernel, iterations=1)
img_rgba[:, :, 3] = alpha


# --- 3. CROP AND SAVE ---
# Crop to content
coords = cv2.findNonZero(img_rgba[:, :, 3])
if coords is not None:
    x, y, w_content, h_content = cv2.boundingRect(coords)

    # Padding
    padding = 2
    x = max(0, x - padding)
    y = max(0, y - padding)
    w_content = min(w - x, w_content + 2*padding)
    h_content = min(h - y, h_content + 2*padding)

    cropped_main = img_rgba[y:y+h_content, x:x+w_content]
else:
    cropped_main = img_rgba # Fallback

# Save FINAL asset v7
output_path = os.path.join(output_dir, 'stream_final_v7.png')
cv2.imwrite(output_path, cropped_main)
print(f"Saved V7 asset to {output_path}")