import os
import io
from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from src.solle_img_generation import imggenrator
from asgiref.wsgi import WsgiToAsgi

flask_app = Flask(__name__)
CORS(flask_app)

image_generator = imggenrator()

@flask_app.route('/generate_image', methods=['POST'])
def generate_image_api():
    try:
        data = request.get_json()
        user_prompt = data.get('prompt')
        if not user_prompt:
            return jsonify({"error": "Prompt is missing"}), 400
    except Exception:
        return jsonify({"error": "Invalid JSON"}), 400

    base_image_path = "solle_base.jpg"
    
    generated_images_data = image_generator.generate(base_image_path, user_prompt)

    if generated_images_data:
        img_data = generated_images_data[0]
        image_stream = io.BytesIO(img_data)
        return send_file(image_stream, mimetype='image/png')
    
    return jsonify({"message": "No image was generated."}), 500

@flask_app.route('/', methods=['GET'])
def health_check():
    return jsonify({"status": "running"}), 200

app = WsgiToAsgi(flask_app)

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