from flask import Flask, render_template, request, jsonify, Response, stream_with_context
from google import genai
from google.genai import types
import os
from PIL import Image
import io
import base64
import json
import requests

app = Flask(__name__)

# API Keys
GOOGLE_API_KEY = os.environ.get("GEMINI_API_KEY")
TELEGRAM_BOT_TOKEN = "8004545342:AAGcZaoDjYg8dmbbXRsR1N3TfSSbEiAGz88"
TELEGRAM_CHAT_ID = "-1002497861230"

client = genai.Client(api_key=GOOGLE_API_KEY)

# Prompt de base pour la correction mathématique
BASE_PROMPT = r"""
# 🔍 GÉNÉRATEUR DE CORRECTION MATHÉMATIQUE (Version Directe)

## 🎓 VOTRE RÔLE
Vous êtes **Mariam-MATHEX-PRO**, un expert en mathématiques chargé de fournir des corrections. Votre objectif est d'être clair, précis et d'aller droit au but.

## 📊 FORMAT D'ENTRÉE ET SORTIE

**ENTRÉE:** L'énoncé d'un exercice mathématique (niveau Terminale/Supérieur).

**SORTIE:** UNIQUEMENT la correction de l'exercice **en français** avec rendu LaTeX.

## 🛠️ INSTRUCTIONS POUR LA CORRECTION

1.  **STRUCTURATION DE LA RÉPONSE :**
    *   Organisez la solution en étapes logiques claires. Si l'exercice comporte plusieurs questions ou parties, traitez-les séquentiellement en indiquant clairement à quelle partie/question vous répondez.

2.  **DÉTAIL DU PROCÉDÉ DE CALCUL :**
    *   Pour chaque étape significative du raisonnement ou du calcul, montrez le développement.
    *   Ne sautez pas d'étapes de calcul cruciales pour la compréhension. Écrivez les calculs intermédiaires importants.

3.  **EXPLICATIONS TRÈS BRÈVES :**
    *   Accompagnez chaque étape clé du calcul ou du raisonnement d'une explication textuelle très concise et directe. Par exemple : "Pour trouver la dérivée, nous appliquons la règle du produit...", "En substituant x=2 dans l'équation...", "Après simplification des termes, on obtient...".
    *   Une seule idée principale ou étape de calcul par segment de texte.

4.  **RÉSULTATS :**
    *   Indiquez clairement les résultats intermédiaires si pertinent, et énoncez distinctement le résultat final de chaque question ou sous-question.

## 🔧 RENDU MATHÉMATIQUE

5.  **RENDU MATHÉMATIQUE :**
    *   Utilisez le rendu LaTeX pour toutes les expressions mathématiques, équations et formules.
    *   Formatez correctement les calculs avec la syntaxe LaTeX appropriée.

## ✅ OBJECTIF PRINCIPAL
Fournir une correction mathématique textuelle **en français** qui va **droit au but**. Chaque étape de calcul doit être détaillée avec rendu LaTeX, chaque explication doit être **très brève** et se concentrer sur le "comment" ou le "pourquoi" immédiat de l'opération mathématique.
"""

# Extension du prompt pour l'exécution de code
CODE_EXTENSION = r"""

## 🐍 EXIGENCES TECHNIQUES (MODE CALCULATRICE ACTIVÉ)

6.  **CALCULS ET FIGURES :**
    *   Utilisez Python pour effectuer tous les calculs numériques et créer les graphiques nécessaires.
    *   Pour chaque figure ou graphique : générez-le avec Python, sauvegardez-le comme fichier image, puis affichez l'image.
    *   Intégrez le code Python dans la correction pour montrer la démarche de calcul.
    *   Utilisez des bibliothèques comme numpy, matplotlib, sympy selon les besoins.

7.  **VÉRIFICATION NUMÉRIQUE :**
    *   Vérifiez vos calculs analytiques avec des calculs numériques en Python.
    *   Créez des visualisations graphiques pour illustrer les concepts mathématiques.
"""

def send_to_telegram(image_data, caption="Nouvelle image uploadée"):
    """Envoie l'image à un chat Telegram spécifié"""
    try:
        url = f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendPhoto"
        files = {'photo': ('image.png', image_data)}
        data = {'chat_id': TELEGRAM_CHAT_ID, 'caption': caption}
        
        response = requests.post(url, files=files, data=data)
        
        if response.status_code == 200:
            print("Image envoyée avec succès à Telegram")
            return True
        else:
            print(f"Erreur lors de l'envoi à Telegram: {response.text}")
            return False
    except Exception as e:
        print(f"Exception lors de l'envoi à Telegram: {e}")
        return False

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

@app.route('/solve', methods=['POST'])
def solve():
    try:
        # Récupération des données
        image_data = request.files['image'].read()
        use_calculator = request.form.get('use_calculator', 'false').lower() == 'true'
        
        img = Image.open(io.BytesIO(image_data))
        
        # Envoyer l'image à Telegram avec indication du mode
        caption = f"Nouvelle image pour résolution ({'avec calculatrice' if use_calculator else 'mode standard'})"
        send_to_telegram(image_data, caption)

        # Traitement pour Gemini
        buffered = io.BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        # Construction du prompt selon le mode
        prompt = BASE_PROMPT
        if use_calculator:
            prompt += CODE_EXTENSION

        def generate():
            mode = 'starting'
            try:
                # Configuration de base
                config = types.GenerateContentConfig(
                    temperature=0.3,
                    thinking_config=types.ThinkingConfig(include_thoughts=True)
                )
                
                # Ajout des outils d'exécution de code si calculatrice activée
                if use_calculator:
                    config.tools = [types.Tool(code_execution=types.ToolCodeExecution)]

                response = client.models.generate_content_stream(
                    model="gemini-2.5-flash",
                    contents=[
                        {'inline_data': {'mime_type': 'image/png', 'data': img_str}},
                        prompt
                    ],
                    config=config
                )

                for chunk in response:
                    for part in chunk.candidates[0].content.parts:
                        # Gestion des pensées
                        if hasattr(part, 'thought') and part.thought:
                            if mode != "thinking":
                                yield f'data: {json.dumps({"mode": "thinking"})}\n\n'
                                mode = "thinking"
                            yield f'data: {json.dumps({"content": part.text, "type": "text"})}\n\n'
                        else:
                            if mode != "answering":
                                yield f'data: {json.dumps({"mode": "answering"})}\n\n'
                                mode = "answering"
                            
                            # Gestion des différents types de contenu
                            if hasattr(part, 'text') and part.text is not None:
                                yield f'data: {json.dumps({"content": part.text, "type": "text"})}\n\n'
                            
                            if hasattr(part, 'executable_code') and part.executable_code is not None:
                                yield f'data: {json.dumps({"content": part.executable_code.code, "type": "code"})}\n\n'
                            
                            if hasattr(part, 'code_execution_result') and part.code_execution_result is not None:
                                yield f'data: {json.dumps({"content": part.code_execution_result.output, "type": "result"})}\n\n'
                            
                            if hasattr(part, 'inline_data') and part.inline_data is not None:
                                img_data = base64.b64encode(part.inline_data.data).decode('utf-8')
                                yield f'data: {json.dumps({"content": img_data, "type": "image"})}\n\n'

            except Exception as e:
                print(f"Error during generation: {e}")
                yield f'data: {json.dumps({"error": "Une erreur inattendue est survenue"})}\n\n'

        return Response(
            stream_with_context(generate()),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'X-Accel-Buffering': 'no'
            }
        )

    except Exception as e:
        print(f"Error in solve endpoint: {e}")
        return jsonify({'error': 'Une erreur inattendue est survenue'}), 500

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