import gradio as gr
import os
import tempfile
from PIL import Image
import subprocess
import logging
import base64
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor, as_completed
import re
from iconsheet import icon_sheet_interface
import mimetypes

# Set up logging
logging.basicConfig(filename='svgmaker.log', level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')

def process_single_image(image_file, stroke_width, fill, stroke, opacity, output_size, output_dir):
    try:
        if image_file is None:
            raise ValueError("No image file provided.")
        
        logging.info(f"Processing image: {image_file.name}")
        with Image.open(image_file.name) as image:
            # Calculate new dimensions maintaining aspect ratio
            aspect_ratio = image.width / image.height
            if aspect_ratio > 1:
                new_width = output_size
                new_height = int(output_size / aspect_ratio)
            else:
                new_height = output_size
                new_width = int(output_size * aspect_ratio)
            
            logging.info(f"Resizing image to {new_width}x{new_height}")
            # Resize image
            image = image.resize((new_width, new_height), Image.LANCZOS)
            
            # Convert to grayscale
            image = image.convert('L')
            
            # Save as temporary PGM file
            with tempfile.NamedTemporaryFile(suffix='.pgm', delete=False) as temp_pgm:
                image.save(temp_pgm.name)
                temp_pgm_path = temp_pgm.name

        # Run potrace with improved settings
        output_svg = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(image_file.name))[0]}.svg")
        subprocess.run([
            "potrace",
            temp_pgm_path,
            "--svg",
            "-o", output_svg            
        ], check=True)

        # Clean up temporary PGM
        os.unlink(temp_pgm_path)

        # Create a base64 encoded version of the SVG for preview
        with open(output_svg, 'r') as f:
            svg_content = f.read()
        svg_base64 = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
        preview = f"data:image/svg+xml;base64,{svg_base64}"

        return output_svg, preview
    except Exception as e:
        logging.error(f"Error processing image {image_file.name}: {str(e)}")
        return None, None

def vectorize_icons(images, stroke_width, fill, stroke, opacity, output_size, output_dir, progress=gr.Progress()):
    if images is None or len(images) == 0:
        raise ValueError("No images were uploaded. Please upload at least one image.")

    vectorized_icons = []
    svg_previews = []
    total_images = len(images)
    logging.info(f"Vectorizing {total_images} images")

    os.makedirs(output_dir, exist_ok=True)

    with ThreadPoolExecutor() as executor:
        future_to_image = {executor.submit(process_single_image, image, stroke_width, fill, stroke, opacity, output_size, output_dir): image for image in images}
        for i, future in enumerate(as_completed(future_to_image)):
            image = future_to_image[future]
            progress((i + 1) / total_images, f"Vectorizing image {i+1}/{total_images}")
            try:
                result, preview = future.result()
                vectorized_icons.append(result)
                svg_previews.append(preview)
                logging.info(f"Successfully vectorized image {i+1}")
                progress((i + 1) / total_images, f"Vectorized image {i+1}/{total_images}")
            except Exception as e:
                logging.error(f"Error vectorizing icon {image.name}: {str(e)}")
                vectorized_icons.append(None)
                svg_previews.append(None)
                progress((i + 1) / total_images, f"Failed to vectorize image {i+1}/{total_images}: {str(e)}")

    progress(1.0, "Vectorization complete")
    logging.info(f"Vectorization complete. Successful: {len([i for i in vectorized_icons if i is not None])}, Failed: {len([i for i in vectorized_icons if i is None])}")
    return vectorized_icons, svg_previews

def create_preview_grid(svg_previews):
    # Dynamically create grid based on number of SVG previews
    num_images = len([preview for preview in svg_previews if preview is not None])
    
    # Define the number of columns for the grid
    columns = 4
    grid_html = f'<div style="display: grid; grid-template-columns: repeat({columns}, 1fr); gap: 10px; background-color: #f0f0f0; padding: 20px; border-radius: 10px;">'
    
    for preview in svg_previews:
        if preview:
            grid_html += f'''
                <div style="background-color: white; padding: 10px; border-radius: 5px;">
                    <img src="{preview}" style="width: 100%; height: auto;" onclick="this.classList.toggle('zoomed')" class="zoomable">
                </div>
            '''
        else:
            grid_html += '<div style="background-color: white; padding: 10px; border-radius: 5px;"><p>Error</p></div>'
    
    grid_html += '</div>'
    
    # Add CSS for zoom functionality
    grid_html += '''
    <style>
        .zoomable { transition: transform 0.3s ease; }
        .zoomable.zoomed { transform: scale(2.5); z-index: 1000; position: relative; }
    </style>
    '''
    
    return grid_html


with gr.Blocks() as app:
    gr.Markdown("# SVG Maker and Icon Sheet Generator")
    
    with gr.Tabs():
        with gr.TabItem("SVG Maker"):
            gr.Markdown("# SVG Maker")
    
            with gr.Row():
                image_input = gr.File(label="Upload Images", file_count="multiple")
            
            with gr.Row():
                stroke_width = gr.Slider(0, 10, label="Stroke Width", value=1)
                fill = gr.ColorPicker(label="Fill Color", value="#000000")
                stroke = gr.ColorPicker(label="Stroke Color", value="#000000")
                opacity = gr.Slider(0, 1, label="Opacity", value=1, step=0.1)
            
            with gr.Row():
                output_size = gr.Slider(32, 1024, label="Output Size (px)", value=512, step=32)
                output_dir = gr.Textbox(label="Output Directory", value="output", placeholder="Enter output directory path")
            
            vectorize_btn = gr.Button("Vectorize")
            svg_output = gr.File(label="Vectorized Icons", file_count="multiple")
            svg_preview = gr.HTML(label="SVG Previews")
            
            def process_and_display(images, stroke_width, fill, stroke, opacity, output_size, output_dir):
                if images is None or len(images) == 0:
                    raise ValueError("No images were uploaded. Please upload at least one image.")
                
                vectorized_icons, svg_previews = vectorize_icons(images, stroke_width, fill, stroke, opacity, output_size, output_dir)
                preview_html = create_preview_grid(svg_previews)
                # Filter out None values from vectorized_icons
                valid_icons = [icon for icon in vectorized_icons if icon is not None]
                return valid_icons, preview_html

            vectorize_btn.click(process_and_display, 
                        inputs=[image_input, stroke_width, fill, stroke, opacity, output_size, output_dir], 
                        outputs=[svg_output, svg_preview])
                
        with gr.TabItem("Icon Sheet"):
            mimetypes.add_type('image/svg+xml', '.svg')
            icon_sheet_interface()

if __name__ == "__main__":
    try:
        app.launch()
    except Exception as e:
        logging.error(f"Failed to launch app: {str(e)}")
        raise