import cairosvg
import gradio as gr
from PIL import Image
import os

import cairosvg
import xml.etree.ElementTree as ET
from PIL import Image

def convert_svg_to_png(svg_file, size_input):
    """Convert SVG to PNG with user-specified dimensions, preserving the entire content."""
    # Input validation
    if svg_file is None or not size_input or not size_input.strip():
        return None, None, "Please upload an SVG file and specify a size."
    
    # Read the SVG content
    with open(svg_file.name, 'rb') as f:
        svg_content = f.read()
    
    # Parse the size input
    try:
        output_width, output_height = map(int, size_input.split('x'))
        if output_width <= 0 or output_height <= 0:
            raise ValueError("Width and height must be positive")
    except ValueError:
        return None, None, "Invalid size format. Use 'width x height' (e.g., '800x600') with positive numbers."
    
    # Parse the SVG to handle viewBox and set width/height if missing
    try:
        root = ET.fromstring(svg_content)
        view_box = root.get('viewBox')
        if view_box and 'width' not in root.attrib and 'height' not in root.attrib:
            _, _, vb_width, vb_height = map(float, view_box.split())
            root.set('width', str(vb_width))
            root.set('height', str(vb_height))
            svg_content = ET.tostring(root, encoding='utf-8')
    except ET.ParseError:
        return None, None, "Invalid SVG file."
    
    # Define output path
    output_path = "./output.png"
    
    # Convert SVG to PNG with user-specified dimensions
    cairosvg.svg2png(
        bytestring=svg_content,
        write_to=output_path,
        output_width=output_width,
        output_height=output_height
    )
    
    # Load the PNG
    final_image = Image.open(output_path)
    return final_image, output_path, f"PNG generated at {output_width}x{output_height} pixels."
    
with gr.Blocks() as bl:
    gr.Markdown("# SVG to PNG Converter")
    
    with gr.Row():
        with gr.Column():
            svg_input = gr.File(label="Upload SVG File", file_types=[".svg"])
            size_input = gr.Textbox(
                label="Output Size (width x height)",
                placeholder="e.g., 800x600"
            )
            convert_btn = gr.Button("Convert")
        
        with gr.Column():
            final_output = gr.Image(type='pil', label="Output PNG", height=800)
            download_btn = gr.File(label="Download PNG")
            status_output = gr.Textbox(label="Status", interactive=False)
    
    # Connect the conversion function
    convert_btn.click(
        fn=convert_svg_to_png,
        inputs=[svg_input, size_input],
        outputs=[final_output, download_btn, status_output]
    )

bl.launch()