import gradio as gr
import cv2
import numpy as np
from PIL import Image
import io
import os
import tempfile

def image_to_line_art_svg(image, line_thickness=2, blur_value=7, threshold_value=7):
    """
    Convert an image to line art and return as SVG
    """
    try:
        # Ensure blur_value and threshold_value are odd
        blur_value = blur_value if blur_value % 2 == 1 else blur_value + 1
        threshold_value = threshold_value if threshold_value % 2 == 1 else threshold_value + 1
        
        # Convert PIL image to numpy array
        if isinstance(image, Image.Image):
            image_array = np.array(image.convert('RGB'))
        else:
            image_array = image
        
        # Resize image if too large (to prevent memory issues)
        height, width = image_array.shape[:2]
        max_dimension = 1024
        if max(height, width) > max_dimension:
            scale = max_dimension / max(height, width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            image_array = cv2.resize(image_array, (new_width, new_height))
            height, width = new_height, new_width
        
        # Convert to grayscale
        gray = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY)
        
        # Apply median filter to reduce noise
        gray_blur = cv2.medianBlur(gray, blur_value)
        
        # Create an edge mask using adaptive threshold
        edges = cv2.adaptiveThreshold(gray_blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C, 
                                     cv2.THRESH_BINARY, threshold_value, blur_value)
        
        # Convert edges back to 3-channel for display
        edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
        
        # Invert edges so lines are black on white background
        edges = cv2.bitwise_not(edges)
        
        # Find contours
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Create SVG content
        svg_content = f'''<?xml version="1.0" encoding="UTF-8"?>
<svg width="{width}" height="{height}" viewBox="0 0 {width} {height}" xmlns="http://www.w3.org/2000/svg">
<style>
    .line-art {{ 
        fill: none; 
        stroke: #000000; 
        stroke-width: {line_thickness}; 
        stroke-linecap: round; 
        stroke-linejoin: round; 
    }}
</style>
<rect width="100%" height="100%" fill="white"/>'''
        
        # Add paths for each contour
        for contour in contours:
            if len(contour) > 5:  # Only include contours with sufficient points
                # Simplify contour to reduce SVG size
                epsilon = 0.01 * cv2.arcLength(contour, True)
                simplified_contour = cv2.approxPolyDP(contour, epsilon, True)
                
                if len(simplified_contour) > 2:
                    path_data = f"M {simplified_contour[0][0][0]} {simplified_contour[0][0][1]}"
                    
                    for point in simplified_contour[1:]:
                        path_data += f" L {point[0][0]} {point[0][1]}"
                    
                    # Close the path if it forms a closed shape
                    if len(simplified_contour) > 3:
                        path_data += " Z"
                    
                    svg_content += f'\n<path d="{path_data}" class="line-art"/>'
        
        svg_content += '\n</svg>'
        
        return svg_content, edges_colored
        
    except Exception as e:
        # Return error SVG
        error_svg = f'''<?xml version="1.0" encoding="UTF-8"?>
<svg width="400" height="200" xmlns="http://www.w3.org/2000/svg">
<rect width="100%" height="100%" fill="white"/>
<text x="200" y="100" text-anchor="middle" fill="red" font-family="Arial" font-size="16">Error: {str(e)}</text>
</svg>'''
        error_image = np.ones((200, 400, 3), dtype=np.uint8) * 255
        return error_svg, error_image

def process_image(image, line_thickness, blur_value, threshold_value):
    """
    Process the uploaded image and return both SVG content and preview image
    """
    if image is None:
        return None, None, "❌ Please upload an image first."
    
    try:
        svg_content, preview_image = image_to_line_art_svg(
            image, line_thickness, blur_value, threshold_value
        )
        
        # Create temporary file for SVG download
        temp_dir = tempfile.gettempdir()
        svg_path = os.path.join(temp_dir, "line_art.svg")
        
        with open(svg_path, "w", encoding="utf-8") as f:
            f.write(svg_content)
        
        return preview_image, svg_path, "✅ Conversion completed successfully!"
        
    except Exception as e:
        error_image = np.ones((200, 400, 3), dtype=np.uint8) * 255
        return error_image, None, f"❌ Error processing image: {str(e)}"

# Create Gradio interface
with gr.Blocks(
    title="Image to Line Art SVG Converter",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1200px !important;
    }
    """
) as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1>🎨 Image to Line Art SVG Converter</h1>
        <p style="font-size: 18px; color: #666;">
            Convert your images into beautiful line art and download as scalable SVG files!
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload & Settings")
            
            image_input = gr.Image(
                label="Upload Image", 
                type="pil",
                sources=["upload"],
                height=300
            )
            
            with gr.Group():
                gr.Markdown("**Adjustment Controls**")
                line_thickness = gr.Slider(
                    minimum=0.5, 
                    maximum=5, 
                    value=2, 
                    step=0.1,
                    label="📏 Line Thickness",
                    info="Controls the width of the drawn lines"
                )
                blur_value = gr.Slider(
                    minimum=3, 
                    maximum=15, 
                    value=7, 
                    step=2,
                    label="🌀 Blur Amount",
                    info="Higher values = smoother lines (must be odd)"
                )
                threshold_value = gr.Slider(
                    minimum=3, 
                    maximum=15, 
                    value=7, 
                    step=2,
                    label="🎯 Edge Sensitivity",
                    info="Lower values = more details (must be odd)"
                )
            
            convert_btn = gr.Button(
                "🚀 Convert to Line Art", 
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 👁️ Preview & Download")
            
            preview_output = gr.Image(
                label="Line Art Preview",
                height=300
            )
            
            status_output = gr.Textbox(
                label="Status", 
                interactive=False,
                show_label=False
            )
            
            download_output = gr.File(
                label="📥 Download SVG File",
                file_count="single"
            )
    
    # Add usage instructions
    with gr.Row():
        gr.Markdown("""
        ### 📝 How to Use:
        1. **Upload** an image using the upload area
        2. **Adjust** the settings to fine-tune the line art (optional)
        3. **Click** "Convert to Line Art" to process
        4. **Download** your SVG file to use anywhere!
        
        ### 💡 Tips:
        - **Line Thickness**: Start with 2, increase for bolder lines
        - **Blur Amount**: Higher values create cleaner, simpler lines  
        - **Edge Sensitivity**: Lower values capture more details
        - SVG files are scalable and perfect for print or web use!
        """)
    
    # Event handlers
    convert_btn.click(
        fn=process_image,
        inputs=[image_input, line_thickness, blur_value, threshold_value],
        outputs=[preview_output, download_output, status_output],
        show_progress=True
    )
    
    # Auto-process when image is uploaded (optional)
    image_input.change(
        fn=lambda img: process_image(img, 2, 7, 7) if img is not None else (None, None, ""),
        inputs=[image_input],
        outputs=[preview_output, download_output, status_output],
        show_progress=False
    )

# Launch settings
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )