import gradio as gr
import vtracer
import os
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io

# Initialize FastAPI app
app = FastAPI()

# Configure CORS to allow requests from Figma plugin
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://www.figma.com"],  # Restrict to Figma's domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Existing vtracer conversion function
def convert_to_vector(
    image,
    colormode="color",
    hierarchical="stacked",
    mode="spline",
    filter_speckle=4,
    color_precision=6,
    layer_difference=16,
    corner_threshold=60,
    length_threshold=4.0,
    max_iterations=10,
    splice_threshold=45,
    path_precision=3
):
    input_path = "temp_input.jpg"
    output_path = "svg_output.svg"

    # Save the input image to a temporary file
    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=int(filter_speckle),
        color_precision=int(color_precision),
        layer_difference=int(layer_difference),
        corner_threshold=int(corner_threshold),
        length_threshold=float(length_threshold),
        max_iterations=int(max_iterations),
        splice_threshold=int(splice_threshold),
        path_precision=int(path_precision)
    )

    # Read the SVG output
    with open(output_path, "r") as f:
        svg_content = f.read()

    # Clean up temporary files
    os.remove(input_path)
    os.remove(output_path)

    return svg_content

# FastAPI endpoint for vector conversion
@app.post("/convert")
async def convert_image(
    file: UploadFile = File(...),
    colormode: str = Form("color"),
    hierarchical: str = Form("stacked"),
    mode: str = Form("spline"),
    filter_speckle: int = Form(4),
    color_precision: int = Form(6),
    layer_difference: int = Form(16),
    corner_threshold: int = Form(60),
    length_threshold: float = Form(4.0),
    max_iterations: int = Form(10),
    splice_threshold: int = Form(45),
    path_precision: int = Form(3)
):
    try:
        # Read the uploaded image
        image_data = await file.read()
        image = Image.open(io.BytesIO(image_data))

        # Convert to SVG
        svg_content = convert_to_vector(
            image,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision
        )

        return JSONResponse(content={"svg": svg_content})
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

# Existing Gradio interface (unchanged)
with gr.Blocks(css=...) as gradio_app:
    # Your existing Gradio code here (omitted for brevity)
    pass

# Mount Gradio app to FastAPI
app.mount("/", gradio_app)

# Note: app.launch() is not needed for Hugging Face Spaces; it will be handled by the platform