# Attempt to import cv2 and install opencv-python-headless if not available
try:
    import cv2
except ImportError:
    import subprocess
    subprocess.check_call(["pip", "install", "opencv-python-headless"])
    import cv2

from PIL import Image
import numpy as np
import gradio as gr

# Function to convert image to sketch
def convert_to_sketch(img, blur_strength, brightness, contrast):
    try:
        # Add debug print
        print(f"Input image shape: {np.array(img).shape}")
        
        # Convert PIL Image to numpy array (BGR format for OpenCV)
        img = np.array(img)
        
        # Check if image is empty
        if img is None or img.size == 0:
            print("Error: Empty input image")
            return None
            
        # Ensure blur_strength is odd
        blur_strength = max(1, int(blur_strength))
        if blur_strength % 2 == 0:
            blur_strength += 1
            
        # Convert the image to grayscale
        img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        print(f"Grayscale shape: {img_gray.shape}")
        
        # Invert the grayscale image
        img_inverted = 255 - img_gray
        
        # Apply Gaussian blur to the inverted image
        img_blur = cv2.GaussianBlur(img_inverted, (blur_strength, blur_strength), 0)
        
        # Blend the grayscale and blurred inverted images
        denominator = 255 - img_blur
        denominator[denominator == 0] = 1  # Prevent division by zero
        
        img_blend = (img_gray * 256.0) / denominator
        
        # Clip values to valid range
        img_blend = np.clip(img_blend, 0, 255).astype(np.uint8)
        print(f"Blend shape: {img_blend.shape}")
        
        # Adjust brightness and contrast
        sketch_with_bg = adjust_brightness_contrast(img_blend, brightness, contrast)
        print(f"Final output shape: {sketch_with_bg.shape}")
        
        return sketch_with_bg
        
    except Exception as e:
        print(f"Error in convert_to_sketch: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# Function to adjust brightness and contrast
def adjust_brightness_contrast(img, brightness, contrast):
    try:
        # Ensure brightness is within valid range
        brightness = float(brightness)
        contrast = float(contrast)
        
        # Convert to float32 for calculations
        img = img.astype(float)
        
        # Apply contrast
        img = img * contrast
        
        # Apply brightness
        img = img + brightness
            
        # Ensure output is in valid range
        img = np.clip(img, 0, 255).astype(np.uint8)
        return img
        
    except Exception as e:
        print(f"Error in adjust_brightness_contrast: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# Gradio interface function
def sketch_interface(image, blur_strength, brightness, contrast):
    if image is None:
        print("Error: No input image provided!")
        return None
        
    try:
        print(f"Input image type: {type(image)}")
        print(f"Input image size: {image.size}")
        
        # Convert the input image to a sketch with adjustments
        sketch = convert_to_sketch(image, blur_strength, brightness, contrast)
        
        if sketch is None:
            print("Error: Sketch conversion failed!")
            return None
            
        # Convert the processed numpy array back to a PIL Image
        output_image = Image.fromarray(sketch)
        print(f"Output image size: {output_image.size}")
        
        return output_image
        
    except Exception as e:
        print(f"Error in sketch_interface: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# Create Gradio interface
interface = gr.Interface(
    fn=sketch_interface,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=1, maximum=51, step=2, value=21, label="Blur Strength"),
        gr.Slider(minimum=-100, maximum=100, value=0, label="Brightness"),
        gr.Slider(minimum=0.1, maximum=3, step=0.1, value=1, label="Contrast")
    ],
    outputs=gr.Image(type="pil", label="Sketch Output"),
    title="Image to coloring page Converter",
    description="Upload an image to convert it into a sketch, adjust the blur strength, brightness, and contrast for different effects."
)

# Launch the Gradio app
interface.launch()