from flask import Flask, request, jsonify, render_template, send_file
from flask_cors import CORS
from PIL import Image
import io
import json
import os
import uuid
import google.generativeai as genai
import cv2
import numpy as np

generation_config = {
    "temperature": 1,
    "max_output_tokens": 8192,
}

safety_settings = [
    {
        "category": "HARM_CATEGORY_HARASSMENT",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_HATE_SPEECH",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        "threshold": "BLOCK_NONE"
    },
]

GOOGLE_API_KEY = os.environ.get("TOKEN")

# Configuration de l'API Gemini
genai.configure(api_key=GOOGLE_API_KEY)

app = Flask(__name__)
CORS(app)

@app.route('/', methods=['GET'])
def svt():
    """Renders the SVT page."""
    return render_template("svt.html")

# Prompt pour la détection d'objets
DETECTION_PROMPT = "Detect items, with no more than 20 items. Output a json list where each entry contains the 2D bounding box in \"box_2d\" and a text label in \"label\"."

# Prompt pour la description d'image satellite militaire
DESCRIPTION_PROMPT = """
Décrivez en détail cette image satellite militaire. Soyez précis et exhaustif dans votre analyse. 
Identifiez les éléments clés tels que :
- **Infrastructures** : Bâtiments, routes, ponts, aéroports, ports, etc.
- **Véhicules** : Chars, avions, navires, véhicules de transport de troupes, etc.
- **Unités militaires** : Formations de troupes, positions d'artillerie, camps, etc.
- **Défenses** : Bunkers, tranchées, barbelés, etc.
- **Éléments géographiques** : Relief, végétation, cours d'eau, etc.
- **Activités** : Mouvements de troupes, entraînements, constructions, etc.
- **Anomalies** : Tout ce qui semble inhabituel ou suspect.

Fournissez une évaluation globale de la situation et des implications stratégiques possibles.
"""

# Dossier pour enregistrer temporairement les images
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

@app.route('/analyze', methods=['POST'])
def analyze_image():
    try:
        if 'file' not in request.files:
            return jsonify({'error': 'No file part'}), 400

        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No selected file'}), 400

        if file:
            unique_filename = str(uuid.uuid4()) + os.path.splitext(file.filename)[1]
            filename = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(filename)

            # Charger l'image avec OpenCV
            image = cv2.imread(filename)
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convertir en RGB pour l'API Gemini

            # Préparer l'image pour l'API Gemini
            image_part = {
                "mime_type": "image/jpeg",
                "data": cv2.imencode('.jpg', image_rgb)[1].tobytes()
            }

            # Détection d'objets et dessin des boîtes (avec gestion d'erreur)
            detection_results = []  # Initialiser une liste vide pour les résultats de détection
            try:
                model = genai.GenerativeModel("gemini-1.5-flash-exp", safety_settings=safety_settings, generation_config=generation_config)
                response = model.generate_content([DETECTION_PROMPT, image_part])
                cleaned_response_text = response.text.replace('\n', '')

                if cleaned_response_text.startswith("```json"):
                    cleaned_response_text = cleaned_response_text[7:]
                if cleaned_response_text.endswith("```"):
                    cleaned_response_text = cleaned_response_text[:-3]

                detection_results = json.loads(cleaned_response_text)

                # Dessiner les boîtes englobantes avec OpenCV
                for item in detection_results:
                    box = item['box_2d']
                    label = str(item['label'])
                    x0, y0, x1, y1 = box
                    cv2.rectangle(image, (x0, y0), (x1, y1), (255, 0, 0), 2)
                    cv2.putText(image, label, (x0, y0 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

            except Exception as e:
                print(f"Erreur lors de la détection d'objets ou du dessin des boîtes : {e}")

            # Enregistrer l'image traitée
            output_filename = os.path.join(app.config['UPLOAD_FOLDER'], 'output_' + unique_filename)
            cv2.imwrite(output_filename, image)

            # Générer la description
            model = genai.GenerativeModel("gemini-2.0-flash-exp", safety_settings=safety_settings, generation_config=generation_config)
            response = model.generate_content([DESCRIPTION_PROMPT, image_part])
            description = response.text

            return jsonify({
                'image_path': '/uploads/' + 'output_' + unique_filename,
                'description': description,
                'detected_objects': detection_results
            })

    except Exception as e:
        print(f"Une erreur s'est produite : {e}")
        return jsonify({'error': f'Erreur lors du traitement de l\'image : {e}'}), 500

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_file(os.path.join(app.config['UPLOAD_FOLDER'], filename))

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