from flask import Flask, render_template, send_from_directory,  send_file, request, jsonify
from flask_cors import CORS
import subprocess
import os
import io
from model.src.utils.arg_parser import eval_parse_args  # Nuovo import corretto
import sys
from PIL import Image
import base64
import traceback


from model.src import eval

app = Flask(__name__)
CORS(app)

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

@app.route('/generate-design', methods=['GET','POST'])
def generate_design():
    try:
        print("### Inizio generate_design ###")  # Stampa per il debug
        
        # Getting json
        json_data_from_req = request.get_json()
        if not json_data_from_req:
            print("### JSON non ricevuto ###")  # Stampa per il debug
            return "Invalid or missing JSON data", 400
        print("### JSON ricevuto:", json_data_from_req)  # Stampa per il debug

        # Getting Image
        if 'image' not in json_data_from_req:
            print("### Nessuna immagine ricevuta ###")  # Stampa per il debug
            return "No image file in request", 400

        encoded_image = json_data_from_req['image']

        try:
            image_data = base64.b64decode(encoded_image)
            image = Image.open(io.BytesIO(image_data))
            print("### Immagine aperta con successo ###")  # Stampa per il debug
        except Exception as e:
            print(f"### Errore nell'apertura dell'immagine: {str(e)} ###")  # Stampa per il debug
            return f"Failed to open the image: {str(e)}", 400

        # Creazione buffer di immagine
        img_sketch_buffer = io.BytesIO()
        image.save(img_sketch_buffer, format='JPEG')
        img_sketch_buffer.seek(0)
        print("### Immagine salvata nel buffer ###")  # Stampa per il debug

        # Argomenti per eval.main
        sys.argv = [
            'eval.py',
            '--dataset_path', '/api/model/assets/data/vitonhd',
            '--batch_size', '1',
            '--mixed_precision', 'fp16',
            '--num_workers_test', '4',
            '--sketch_cond_rate', '0.2',
            '--dataset', 'vitonhd',
            '--start_cond_rate', '0.0',
            '--test_order', 'paired'
        ]
        print("### Esecuzione eval.main ###")  # Stampa per il debug

        import traceback

        # Esecuzione del modello
        try:
            final_image = eval.main(img_sketch_buffer, json_data_from_req)
            print("### eval.main eseguito con successo ###")  # Stampa per il debug
        except AttributeError as e:
            print(f"### AttributeError: {str(e)} ###")  # Stampa per il debug
            return f"AttributeError: {traceback.format_exc()}", 500
        except TypeError as e:
            print(f"### TypeError: {str(e)} ###")  # Stampa per il debug
            return f"TypeError: {traceback.format_exc()}", 500
        except ValueError as e:
            print(f"### ValueError: {str(e)} ###")  # Stampa per il debug
            return f"ValueError: {traceback.format_exc()}", 500
        except IOError as e:
            print(f"### IOError: {str(e)} ###")  # Stampa per il debug
            return f"IOError: {traceback.format_exc()}", 500
        except Exception as e:
            print(f"### Unexpected error: {str(e)} ###")  # Stampa per il debug
            return f"Unexpected error: {traceback.format_exc()}", 500


        # Salvataggio immagine e invio come risposta
        img_io = io.BytesIO()
        final_image.save(img_io, 'JPEG')
        img_io.seek(0)
        print("### Immagine finale pronta per il download ###")  # Stampa per il debug

        return send_file(img_io, mimetype='image/jpeg', as_attachment=True, download_name='inmemory_image.jpg')

    except Exception as e:
        print(f"### Errore globale: {str(e)} ###")  # Stampa per il debug
        return str(e), 500
    
@app.route('/test_post', methods=['POST'])
def test_post():
    # Verifica se la richiesta è POST
    if request.method == 'POST':
        # Ottieni i dati dal corpo della richiesta
        data = request.get_json()
        
        # Se i dati non sono presenti o il formato non è corretto
        if not data:
            return jsonify({'message': 'Nessun dato ricevuto o dati non validi'}), 400
        
        # Restituisci una risposta con i dati ricevuti
        return jsonify({'message': 'POST funzionante', 'data_received': data}), 200

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