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

# Configuration de l'API Gemini
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") # Assurez vous que la variable d'environnement TOKEN est bien définie

genai.configure(api_key=GOOGLE_API_KEY)

app = Flask(__name__)
CORS(app)

# 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('/', methods=['GET'])
def svt():
    """Renders the SVT page."""
    return render_template("svt.html")

@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:
            # Générer un nom de fichier unique pour éviter les conflits
            unique_filename = str(uuid.uuid4()) + os.path.splitext(file.filename)[1]
            filename = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(filename)

            # 1. Détection d'objets avec Gemini
            model = genai.GenerativeModel("gemini-2.0-flash-exp",safety_settings=safety_settings,generation_config=generation_config)
            image_part = {
              "mime_type": "image/jpeg",
              "data": open(filename, "rb").read()
            }
            response = model.generate_content([DETECTION_PROMPT, image_part])

            # Nettoyer la réponse JSON en supprimant les sauts de ligne
            cleaned_response_text = response.text.replace('\n', '')

            try:
                # Enlever les ```json et ``` au début et à la fin de la chaîne
                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)
            except json.JSONDecodeError:
                print(f"Erreur de décodage JSON : {cleaned_response_text}")
                detection_results = []  # Initialiser à une liste vide en cas d'erreur

            # 2. Dessiner les boîtes englobantes (avec gestion d'erreur)
            image = Image.open(filename)
            draw = ImageDraw.Draw(image)

            draw_success = True  # Indicateur de succès du dessin

            for item in detection_results:
                try:
                    box = item['box_2d']
                    label = item['label']

                    # Convertir la liste 'box' en tuple
                    box_tuple = tuple(box)

                    draw.rectangle(box_tuple, outline=(255, 0, 0), width=2)
                    text_position = (box[0], box[1] - 10)

                    # S'assurer que le label est une chaîne de caractère
                    label_str = str(label)

                    # Utiliser une couleur valide (blanc)
                    draw.text(text_position, label_str, fill="white")

                except Exception as e:
                    print(f"Erreur lors du dessin des boîtes ou du texte : {e}")
                    draw_success = False  # Échec du dessin
                    break  # Sortir de la boucle en cas d'erreur

            # 3. Générer la description (toujours exécuté)
            response = model.generate_content([DESCRIPTION_PROMPT, image_part])
            description = response.text

            # 4. Renvoyer les résultats
            if draw_success:
                # Enregistrer l'image avec les boîtes si le dessin a réussi
                output_filename = os.path.join(app.config['UPLOAD_FOLDER'], 'output_' + unique_filename)
                image.save(output_filename)
                return jsonify({
                    'image_path': '/uploads/' + 'output_' + unique_filename,
                    'description': description,
                    'detected_objects': detection_results
                })
            else:
                # Renvoyer la description et les résultats de détection, même si le dessin a échoué
                return jsonify({
                    'image_path': None,  # Indiquer qu'aucune image n'a été générée
                    '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

# Servir les fichiers statiques depuis le dossier 'uploads'
@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)