from flask import Flask, render_template, request, send_file, jsonify
from imaginepy import AsyncImagine #, Style, Ratio, Model
from imaginepy.constants import *
import os

app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
async def generate_image():
    prompt = request.form['prompt']
    style = request.form['style']
    ratio = request.form['ratio']
    model = request.form['model']

    imagine = AsyncImagine()

    try:
        img_data = await imagine.sdprem(
            prompt=prompt,
            style=Style[style],
            ratio=Ratio[ratio],
            seed="",
            cfg=16,
            model=Model[model],
            asbase64=False
        )
    except Exception as e:
        return f"An error occurred while generating the image: {e}"

    if img_data is None:
        return "An error occurred while generating the image."

    img_data = await imagine.upscale(img_data)

    if img_data is None:
        return "An error occurred while upscaling the image."

    try:
        with open("static/example.jpeg", mode="wb") as img_file:
            img_file.write(img_data)
    except Exception as e:
        return f"An error occurred while writing the image to file: {e}"
    finally:
        await imagine.close()
    return render_template('output.html')


@app.route('/api/generate', methods=['POST'])
def api_generate_image():
    data = request.get_json()
    prompt = data['prompt']
    style = data['style']
    ratio = data['ratio']
    model = data['model']

    imagine = Imagine()

    try:
        img_data = imagine.sdprem(
            prompt=prompt,
            style=Style[style],
            ratio=Ratio[ratio],
            seed='',
            cfg=16,
            model=Model[model],
            asbase64=False
        )
    except Exception as e:
        return jsonify({'error': f"An error occurred while generating the image: {e}"}), 500

    if img_data is None:
        return jsonify({'error': "An error occurred while generating the image."}), 500

    img_data = imagine.upscale(img_data)

    if img_data is None:
        return jsonify({'error': "An error occurred while upscaling the image."}), 500

    try:
        image_path = os.path.join(app.root_path, "generated.jpeg")
        with open(image_path, mode="wb") as img_file:
            img_file.write(img_data)
    except Exception as e:
        return jsonify({'error': f"An error occurred while writing the image to file: {e}"}), 500

    finally:
        imagine.close()

    return send_file(image_path, mimetype='image/jpeg', as_attachment=True)
    
@app.errorhandler(404)
def page_not_found(e):
    # Render the 404.html template
    return render_template('404.html'), 404

@app.route('/api-docs')
def api_docs():
    return render_template('api_docs.html')
    
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860,debug=True)