from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from PIL import Image
import vtracer
import io
import os

app = Flask(__name__)
CORS(app, resources={r"/convert-to-vector": {"origins": "*"}})

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 = "temp_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

@app.route('/convert-to-vector', methods=['POST'])
def convert_to_vector_endpoint():
    try:
        # Handle raw image data in the request body
        if request.data:
            image = Image.open(io.BytesIO(request.data)).convert('RGB')
        else:
            return jsonify({'error': 'No image data provided'}), 400

        # Convert image to SVG with default parameters (can be extended with query params if needed)
        svg_content = convert_to_vector(image)

        # Save SVG content to a bytes buffer
        svg_io = io.StringIO(svg_content)
        svg_bytes = io.BytesIO(svg_content.encode('utf-8'))
        svg_bytes.seek(0)

        # Return the SVG as a file response
        return send_file(svg_bytes, mimetype='image/svg+xml', as_attachment=True, download_name='vector_output.svg')
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/')
def index():
    return jsonify({'status': 'Image to Vector Converter API is running'})

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