import gradio as gr
import vtracer
import os
import base64
from io import BytesIO
from PIL import Image
import json
from flask import Flask, request, jsonify
from flask_cors import CORS

app = Flask(__name__)
CORS(app)

@app.route('/api/predict', methods=['POST'])
def predict():
    try:
        # Get image file from request
        if 'image' not in request.files:
            return jsonify({"error": "No image file provided"}), 400
        
        image_file = request.files['image']
        image = Image.open(image_file)
        
        # Get parameters from request
        colormode = request.form.get('colormode', 'color')
        hierarchical = request.form.get('hierarchical', 'stacked')
        mode = request.form.get('mode', 'spline')
        filter_speckle = int(request.form.get('filter_speckle', 4))
        color_precision = int(request.form.get('color_precision', 6))
        layer_difference = int(request.form.get('layer_difference', 16))
        corner_threshold = int(request.form.get('corner_threshold', 60))
        length_threshold = float(request.form.get('length_threshold', 4.0))
        max_iterations = int(request.form.get('max_iterations', 10))
        splice_threshold = int(request.form.get('splice_threshold', 45))
        path_precision = int(request.form.get('path_precision', 3))
        
        # Create temp directory if it doesn't exist
        os.makedirs('temp', exist_ok=True)
        
        # Save the input image to a temporary file
        input_path = "temp/temp_input.jpg"
        output_path = "temp/svg_output.svg"
        
        image.save(input_path)
        
        # Convert the image to SVG using VTracer
        vtracer.convert_image_to_svg_py(
            input_path,
            output_path,
            colormode=colormode,
            hierarchical=hierarchical,
            mode=mode,
            filter_speckle=filter_speckle,
            color_precision=color_precision,
            layer_difference=layer_difference,
            corner_threshold=corner_threshold,
            length_threshold=length_threshold,
            max_iterations=max_iterations,
            splice_threshold=splice_threshold,
            path_precision=path_precision
        )
        
        # Read the SVG output
        with open(output_path, "r") as f:
            svg_content = f.read()
        
        # Create SVG with viewBox
        full_svg = f'<svg viewBox="0 0 {image.width} {image.height}">{svg_content}</svg>'
        
        # Return SVG content and path
        return jsonify({
            "data": [full_svg, output_path],
            "success": True
        })
    
    except Exception as e:
        return jsonify({"error": str(e), "success": False}), 500

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=7860)