import os
from flask import Flask, request, render_template, redirect, url_for
from PIL import Image
import google.generativeai as genai

# Replace with your actual API key
GOOGLE_API_KEY = 'AIzaSyDDt5EWkmYlQTCJZfwNPizs0qZDFznSKVU'

# Configure the API key
genai.configure(api_key=GOOGLE_API_KEY)

# Select the model you want to use
model = genai.GenerativeModel('gemini-1.5-flash-latest')

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'

# Ensure the upload folder exists
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

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

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return redirect(request.url)
    
    file = request.files['file']
    if file.filename == '':
        return redirect(request.url)
    
    if file:
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        file.save(filepath)

        img = Image.open(filepath)

        # Generate content using the model
        response = model.generate_content(["Give the equivalent XML for the UML diagram also take into account all attributes, classes, and functions", img], stream=True)
        response.resolve()

        return render_template('result.html', result=response.text)

if __name__ == '__main__':
    app.run(debug=True)