from flask import Flask, request, send_file, jsonify, render_template, send_from_directory, Response
import io
import base64
import random
import torch
from models.inference import image_gen

app = Flask(__name__)

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

@app.route("/image_gen", methods=['POST'])
def gen():
    image_path = None
    image_bytes = None
    prompts_out = request.form['prompts_out']
    neg_prompts_out = request.form['neg_prompts_out']
    steps_out = int(request.form['steps_out'])
    cfg_out = int(request.form['cfg_out'])
    num_images_out = int(request.form['num_images_out'])
    height_out = int(request.form['height_out'])
    width_out = int(request.form['width_out'])
    clip_skip_out = int(request.form['clip_skip_out'])

    print("---------------------")
    print("prompts out ", prompts_out)
    print("neg prompt", neg_prompts_out)
    print("---------------------")

    image = image_gen(prompts_out, neg_prompts_out, steps_out, cfg_out, num_images_out,
                    height_out, width_out, clip_skip_out)

    # print(image)
    # print("---------------------")
    
    image_bytes = io.BytesIO()
    image.save(image_bytes, format="JPEG", quality=67)
    
    image_bytes = base64.b64encode(image_bytes.getvalue()).decode('utf-8')
    print(f"Generated image of size {len(image_bytes)} bytes.")
        
    return render_template('index.html', image_path=image_bytes)
    

'''
@app.route('/', methods=['POST'])
def myapp():
    # prompts_out = request.args["Prompt"]
    # neg_prompts_out = request.args["Negative Prompt"]
    # steps_out = request.args["Steps"]
    # cfg_out = request.args["CFG"]
    # num_images_out = request.args["Number of Images"]
    # height_out = request.args["Height"]
    # width_out = request.args["Width"]
    # clip_skip_out = request.args["Clip Skip"]
    # img_data = image_gen(prompts_out, neg_prompts_out, steps_out, cfg_out, num_images_out,
    #                     height_out, width_out, clip_skip_out)
    # return send_file(img_data, mimetype='image/png')
    return "Hello World!"
'''


'''
@app.route('/', methods=['GET', 'POST'])
def index():
    image_path = None
    image_bytes =None
    if request.method == 'POST':
        prompts_out = request.form['prompts_out']
        neg_prompts_out = request.form['neg_prompts_out']
        steps_out = int(request.form['steps_out'])
        cfg_out = int(request.form['cfg_out'])
        num_images_out = int(request.form['num_images_out'])
        height_out = int(request.form['height_out'])
        width_out = int(request.form['width_out'])
        clip_skip_out = int(request.form['clip_skip_out'])

        ## seed = random.randrange(1, 100)

        print("---------------------")
        print("prompts out ", prompts_out)
        print("neg prompt", neg_prompts_out)
        print("---------------------")
        
        image = image_gen(prompts_out, neg_prompts_out, steps_out, cfg_out, num_images_out,
                        height_out, width_out, clip_skip_out)

        print(image)
        print("---------------------")
        
        image_bytes = io.BytesIO()
        image.save(image_bytes, format="JPEG", quality=67)
        
        image_bytes = base64.b64encode(image_bytes.getvalue()).decode('utf-8')
        print(f"Generated image of size {len(image_bytes)} bytes.")
        
        # img_str = base64.b64encode(image_bytes.getvalue())
        # img_obj = "data:image/png;base64," + str(img_str)[2:-1]
        
        # return Response(image_bytes, mimetype="image/jpeg")
        
        # return render_template('result.html',image=img_obj)
        
    return render_template('index.html', image_path=image_bytes)
'''

"""
@app.route('/')
def hello():
    return {"test": "Hello world!"}
"""