import os
import json
import uuid
import threading
import time
from datetime import datetime
from flask import Flask, request, render_template, jsonify, send_file
from flask_socketio import SocketIO, emit
from werkzeug.utils import secure_filename
import base64
from io import BytesIO
from PIL import Image
import pytesseract
from crewai import Agent, Task, Crew, Process, LLM
import re
import logging

# Configuration Flask
app = Flask(__name__)
app.config['SECRET_KEY'] = 'votre_cle_secrete_ici'
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max

# Configuration SocketIO pour les logs en temps réel
socketio = SocketIO(app, cors_allowed_origins="*")

# Créer les dossiers nécessaires
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs('solutions', exist_ok=True)

# Extensions d'images autorisées
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff'}

# Configuration Gemini API
os.environ["GEMINI_API_KEY"] = os.environ.get("GEMINI_API_KEY", "AIzaSyAMYpF67aqFnWDJESWOx1dC-w3sEU29VcM")

# Dictionnaire global pour stocker les tâches en cours
active_tasks = {}

# Configuration du LLM
llm = LLM(
    model="gemini/gemini-2.5-flash",
    temperature=0.1,
    
)

# Paramètres du pipeline
MAX_ITERATIONS = 5
PASSES_NEEDED_FOR_SUCCESS = 2

# ==============================================================================
# UTILITAIRES
# ==============================================================================

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def extract_text_from_image(image_path):
    """Extrait le texte d'une image en utilisant OCR."""
    try:
        # Ouvrir l'image
        image = Image.open(image_path)
        
        # Convertir en RGB si nécessaire
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Extraire le texte avec pytesseract
        text = pytesseract.image_to_string(image, lang='eng')
        
        return text.strip()
    except Exception as e:
        return f"Erreur lors de l'extraction du texte: {str(e)}"

def emit_log(task_id, level, message):
    """Émet un log vers le client via WebSocket."""
    log_data = {
        'timestamp': datetime.now().strftime('%H:%M:%S'),
        'level': level,
        'message': message
    }
    socketio.emit('log_update', {'task_id': task_id, 'log': log_data})

# ==============================================================================
# PROMPTS (MÊME LOGIQUE QUE L'ORIGINAL)
# ==============================================================================

def get_imo_paper_solver_prompt(problem_statement, hint=""):
    full_problem = f"{problem_statement}\n{hint}" if hint else problem_statement
    return f"""
### Core Instructions ###
**Rigor is Paramount:** Your primary goal is to produce a complete and rigorously justified solution.
You are solving an International Mathematical Olympiad (IMO) level problem.

### Problem ###
{full_problem}

### Your Task ###
Provide a complete, rigorous mathematical proof. Structure your solution clearly with:
1. A brief summary of your approach
2. A detailed step-by-step proof
3. Clear justification for each step
"""

def get_self_improve_prompt(solution_attempt):
    return f"""
You are a world-class mathematician. You have just produced the following draft solution.
Review it carefully for flaws, gaps, or unclear parts, then produce a new, improved, and more rigorous version.

### Draft Solution ###
{solution_attempt}

### Your Task ###
Provide the improved and finalized version of the solution. Only output the final, clean proof.
"""

def get_imo_paper_verifier_prompt(problem_statement, solution_to_verify):
    return f"""
You are an expert mathematician and a meticulous grader for an IMO level exam.

### Problem ###
{problem_statement}

### Solution ###
{solution_to_verify}

### Verification Task ###
Act as an IMO grader. Analyze the solution for:
1. Mathematical correctness
2. Logical gaps or errors
3. Clarity and rigor

Provide your verdict in the format:
**Final Verdict:** [The solution is correct / Critical error found / Justification gaps found]

Then provide detailed feedback on any issues found.
"""

def get_correction_prompt(solution_attempt, verification_report):
    return f"""
You are a brilliant mathematician. Your previous solution attempt has been reviewed by a verifier.
Your task is to write a new, corrected version of your solution that addresses all the issues raised.

### Verification Report ###
{verification_report}

### Your Previous Solution ###
{solution_attempt}

### Your Task ###
Provide a new, complete, and rigorously correct solution that fixes all identified issues.
"""

# ==============================================================================
# AGENTS CREWAI
# ==============================================================================

solver_agent = Agent(
    role='Mathématicien de classe mondiale et résolveur de problèmes créatifs',
    goal='Générer et affiner des preuves mathématiques complètes et rigoureuses pour des problèmes de niveau Olympiades.',
    backstory="""Médaillé d'or des OIM, vous êtes connu pour votre capacité à trouver des solutions
    à la fois non conventionnelles et impeccablement rigoureuses. Vous excellez à décomposer
    des problèmes complexes et à présenter des arguments clairs, étape par étape.""",
    verbose=True,
    allow_delegation=False,
    llm=llm
)

verifier_agent = Agent(
    role='Examinateur méticuleux pour les Olympiades Internationales de Mathématiques',
    goal='Analyser de manière critique une preuve mathématique pour y trouver la moindre erreur logique ou le moindre manque de justification.',
    backstory="""Professeur de mathématiques réputé, membre du comité de notation des OIM. Votre devise est 'la rigueur avant tout'.
    Vous ne corrigez jamais, vous ne faites que juger et rapporter les failles.""",
    verbose=True,
    allow_delegation=False,
    llm=llm
)

# ==============================================================================
# PIPELINE DE RÉSOLUTION
# ==============================================================================

def parse_verifier_verdict(report):
    if not report: 
        return "ERROR"
    
    report_text = str(report)
    match = re.search(r"\*\*Final Verdict:\*\* (.*)", report_text, re.IGNORECASE)
    if not match: 
        return "UNKNOWN"
    
    verdict_text = match.group(1).lower()
    if "the solution is correct" in verdict_text:
        return "CORRECT"
    if "critical error" in verdict_text:
        return "CRITICAL_ERROR"
    if "justification gap" in verdict_text:
        return "GAPS"
    return "UNKNOWN"

def run_imo_pipeline_async(task_id, problem_statement, initial_hint=""):
    """Version asynchrone du pipeline avec émission de logs."""
    try:
        emit_log(task_id, 'info', "🚀 Démarrage du Pipeline de Résolution Mathématique")
        
        # Étape 1: Génération Initiale
        emit_log(task_id, 'info', "--- ÉTAPE 1: Génération de la Solution Initiale ---")
        
        initial_task = Task(
            description=get_imo_paper_solver_prompt(problem_statement, initial_hint),
            expected_output="Une ébauche de preuve mathématique structurée avec un résumé et une solution détaillée.",
            agent=solver_agent
        )
        
        initial_crew = Crew(
            agents=[solver_agent],
            tasks=[initial_task],
            process=Process.sequential,
            verbose=False
        )
        
        initial_result = initial_crew.kickoff()
        initial_solution = initial_result.raw if hasattr(initial_result, 'raw') else str(initial_result)
        
        emit_log(task_id, 'success', "✅ Solution initiale générée")
        
        # Étape 2: Auto-Amélioration
        emit_log(task_id, 'info', "--- ÉTAPE 2: Auto-Amélioration de la Solution ---")
        
        improve_task = Task(
            description=get_self_improve_prompt(initial_solution),
            expected_output="Une version améliorée et plus rigoureuse de la preuve.",
            agent=solver_agent
        )
        
        improve_crew = Crew(
            agents=[solver_agent],
            tasks=[improve_task],
            process=Process.sequential,
            verbose=False
        )
        
        improve_result = improve_crew.kickoff()
        current_solution = improve_result.raw if hasattr(improve_result, 'raw') else str(improve_result)
        
        emit_log(task_id, 'success', "✅ Solution améliorée")
        
        # Boucle de Vérification
        iteration = 0
        consecutive_passes = 0
        
        while iteration < MAX_ITERATIONS:
            iteration += 1
            emit_log(task_id, 'info', f"--- CYCLE DE VÉRIFICATION {iteration}/{MAX_ITERATIONS} ---")

            verify_task = Task(
                description=get_imo_paper_verifier_prompt(problem_statement, current_solution),
                expected_output="Un rapport de vérification complet au format OIM.",
                agent=verifier_agent
            )
            
            correct_task = Task(
                description=get_correction_prompt(current_solution, "Rapport du vérificateur (voir contexte)"),
                expected_output="Une nouvelle version complète et corrigée de la preuve mathématique.",
                agent=solver_agent,
                context=[verify_task]
            )
            
            correction_crew = Crew(
                agents=[verifier_agent, solver_agent],
                tasks=[verify_task, correct_task],
                process=Process.sequential,
                verbose=False
            )
            
            cycle_result = correction_crew.kickoff()
            
            if hasattr(cycle_result, 'tasks_output') and len(cycle_result.tasks_output) >= 2:
                verification_report = cycle_result.tasks_output[0].raw
                corrected_solution = cycle_result.tasks_output[1].raw
            else:
                verification_report = str(cycle_result)
                corrected_solution = str(cycle_result)
            
            verdict = parse_verifier_verdict(verification_report)
            
            emit_log(task_id, 'info', f"VERDICT DU CYCLE {iteration}: {verdict}")
            
            if verdict == "CORRECT":
                consecutive_passes += 1
                emit_log(task_id, 'success', f"✅ PASSAGE RÉUSSI ! Passes consécutives : {consecutive_passes}/{PASSES_NEEDED_FOR_SUCCESS}")
                if consecutive_passes >= PASSES_NEEDED_FOR_SUCCESS:
                    emit_log(task_id, 'success', "🏆 La solution a été validée avec succès ! 🏆")
                    
                    # Sauvegarder la solution
                    solution_file = f"solutions/solution_{task_id}.txt"
                    with open(solution_file, 'w', encoding='utf-8') as f:
                        f.write(f"Problème:\n{problem_statement}\n\n")
                        f.write(f"Solution finale:\n{current_solution}")
                    
                    active_tasks[task_id]['status'] = 'completed'
                    active_tasks[task_id]['solution_file'] = solution_file
                    
                    socketio.emit('task_completed', {'task_id': task_id})
                    return current_solution
            else:
                consecutive_passes = 0
                emit_log(task_id, 'warning', "⚠️ Problèmes détectés. Correction en cours...")
                current_solution = corrected_solution

        emit_log(task_id, 'error', f"❌ Échec après {MAX_ITERATIONS} itérations")
        
        # Sauvegarder la solution partielle
        solution_file = f"solutions/solution_{task_id}_partial.txt"
        with open(solution_file, 'w', encoding='utf-8') as f:
            f.write(f"Problème:\n{problem_statement}\n\n")
            f.write(f"Solution partielle (non validée):\n{current_solution}")
        
        active_tasks[task_id]['status'] = 'failed'
        active_tasks[task_id]['solution_file'] = solution_file
        
        socketio.emit('task_failed', {'task_id': task_id})
        return current_solution
        
    except Exception as e:
        emit_log(task_id, 'error', f"❌ Erreur critique: {str(e)}")
        active_tasks[task_id]['status'] = 'error'
        active_tasks[task_id]['error'] = str(e)
        socketio.emit('task_error', {'task_id': task_id, 'error': str(e)})

# ==============================================================================
# ROUTES FLASK
# ==============================================================================

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

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'Aucun fichier fourni'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'Aucun fichier sélectionné'}), 400
    
    if file and allowed_file(file.filename):
        # Générer un ID de tâche unique
        task_id = str(uuid.uuid4())
        
        # Sauvegarder le fichier
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{task_id}_{filename}")
        file.save(filepath)
        
        # Extraire le texte de l'image
        try:
            extracted_text = extract_text_from_image(filepath)
            
            # Initialiser la tâche
            active_tasks[task_id] = {
                'status': 'processing',
                'problem': extracted_text,
                'started_at': datetime.now(),
                'filename': filename
            }
            
            # Lancer le pipeline en arrière-plan
            thread = threading.Thread(
                target=run_imo_pipeline_async,
                args=(task_id, extracted_text)
            )
            thread.daemon = True
            thread.start()
            
            return jsonify({
                'task_id': task_id,
                'extracted_text': extracted_text,
                'message': 'Traitement commencé'
            })
            
        except Exception as e:
            return jsonify({'error': f'Erreur lors du traitement de l\'image: {str(e)}'}), 500
    
    return jsonify({'error': 'Type de fichier non autorisé'}), 400

@app.route('/status/<task_id>')
def get_status(task_id):
    if task_id not in active_tasks:
        return jsonify({'error': 'Tâche non trouvée'}), 404
    
    task = active_tasks[task_id]
    return jsonify({
        'task_id': task_id,
        'status': task['status'],
        'problem': task.get('problem', ''),
        'started_at': task['started_at'].isoformat(),
        'has_solution': 'solution_file' in task
    })

@app.route('/download/<task_id>')
def download_solution(task_id):
    if task_id not in active_tasks:
        return jsonify({'error': 'Tâche non trouvée'}), 404
    
    task = active_tasks[task_id]
    if 'solution_file' not in task:
        return jsonify({'error': 'Aucune solution disponible'}), 404
    
    return send_file(task['solution_file'], as_attachment=True, download_name=f'solution_{task_id}.txt')

@app.route('/tasks')
def list_tasks():
    tasks_info = []
    for task_id, task in active_tasks.items():
        tasks_info.append({
            'task_id': task_id,
            'status': task['status'],
            'filename': task.get('filename', ''),
            'started_at': task['started_at'].isoformat(),
            'has_solution': 'solution_file' in task
        })
    return jsonify(tasks_info)

# ==============================================================================
# TEMPLATE HTML
# ==============================================================================

@app.route('/template')
def get_template():
    template_content = '''
    <html lang="fr">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Résolveur Mathématique IMO</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.0.4/socket.io.js"></script>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body { 
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; 
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }
        .container { 
            max-width: 1200px; 
            margin: 0 auto; 
            background: white; 
            border-radius: 15px;
            box-shadow: 0 20px 40px rgba(0,0,0,0.1);
            overflow: hidden;
        }
        .header {
            background: linear-gradient(135deg, #2196F3, #21CBF3);
            padding: 30px;
            text-align: center;
            color: white;
        }
        .header h1 { font-size: 2.5em; margin-bottom: 10px; }
        .header p { font-size: 1.1em; opacity: 0.9; }
        
        .upload-section {
            padding: 40px;
            text-align: center;
            border-bottom: 1px solid #eee;
        }
        .upload-box {
            border: 3px dashed #ddd;
            border-radius: 10px;
            padding: 40px;
            transition: all 0.3s ease;
            cursor: pointer;
        }
        .upload-box:hover { border-color: #2196F3; background: #f8f9ff; }
        .upload-box.dragover { border-color: #2196F3; background: #e3f2fd; }
        
        .btn {
            background: linear-gradient(135deg, #2196F3, #21CBF3);
            color: white;
            padding: 12px 30px;
            border: none;
            border-radius: 25px;
            font-size: 16px;
            cursor: pointer;
            transition: all 0.3s ease;
        }
        .btn:hover { transform: translateY(-2px); box-shadow: 0 5px 15px rgba(33,150,243,0.3); }
        
        .logs-section {
            display: flex;
            height: 600px;
        }
        .extracted-text {
            flex: 1;
            padding: 20px;
            border-right: 1px solid #eee;
        }
        .logs-panel {
            flex: 1;
            padding: 20px;
        }
        .log-container {
            height: 500px;
            overflow-y: auto;
            background: #1e1e1e;
            color: #fff;
            padding: 15px;
            border-radius: 8px;
            font-family: 'Courier New', monospace;
            font-size: 13px;
        }
        .log-entry {
            margin-bottom: 5px;
            padding: 5px;
            border-radius: 3px;
        }
        .log-info { color: #4CAF50; }
        .log-success { color: #8BC34A; background: rgba(139,195,74,0.1); }
        .log-warning { color: #FF9800; }
        .log-error { color: #F44336; background: rgba(244,67,54,0.1); }
        
        .status-bar {
            padding: 20px;
            background: #f5f5f5;
            text-align: center;
        }
        .status-badge {
            display: inline-block;
            padding: 8px 16px;
            border-radius: 20px;
            font-weight: bold;
            text-transform: uppercase;
        }
        .status-processing { background: #FFC107; color: #333; }
        .status-completed { background: #4CAF50; color: white; }
        .status-failed { background: #F44336; color: white; }
        
        .download-section {
            padding: 20px;
            text-align: center;
            display: none;
        }
        .hidden { display: none; }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🧮 Résolveur Mathématique IMO</h1>
            <p>Uploadez une image de votre problème mathématique et obtenez une solution rigoureuse</p>
        </div>
        
        <div class="upload-section">
            <div class="upload-box" id="uploadBox">
                <h3>📁 Glisser-déposer votre image ici</h3>
                <p>ou cliquez pour sélectionner un fichier</p>
                <input type="file" id="fileInput" accept="image/*" style="display: none;">
                <!-- MODIFICATION 1: Suppression du onclick et ajout d'un id -->
                <button class="btn" id="uploadButton">
                    Choisir un fichier
                </button>
            </div>
        </div>
        
        <div class="status-bar">
            <div id="statusBadge" class="status-badge" style="display: none;">En attente</div>
            <div id="taskInfo" style="margin-top: 10px; display: none;"></div>
        </div>
        
        <div class="logs-section hidden" id="logsSection">
            <div class="extracted-text">
                <h3>📝 Texte extrait de l'image</h3>
                <div id="extractedText" style="background: #f9f9f9; padding: 15px; border-radius: 8px; margin-top: 10px; white-space: pre-wrap; max-height: 450px; overflow-y: auto;"></div>
            </div>
            <div class="logs-panel">
                <h3>📊 Logs de traitement en temps réel</h3>
                <div id="logContainer" class="log-container"></div>
            </div>
        </div>
        
        <div class="download-section" id="downloadSection">
            <h3>✅ Solution prête !</h3>
            <button class="btn" id="downloadBtn">📥 Télécharger la solution</button>
        </div>
    </div>
    <script>
        const socket = io();
        let currentTaskId = null;
        
        // Gestion de l'upload
        const uploadBox = document.getElementById('uploadBox');
        const fileInput = document.getElementById('fileInput');
        // MODIFICATION 2: Sélection du bouton par son nouvel id
        const uploadButton = document.getElementById('uploadButton');
        
        // L'utilisateur peut cliquer sur toute la zone
        uploadBox.addEventListener('click', () => fileInput.click());
        
        // MODIFICATION 3: Ajout d'un écouteur dédié au bouton
        uploadButton.addEventListener('click', (e) => {
            // Empêche le clic de "remonter" à la div parente (uploadBox),
            // ce qui évite de déclencher l'ouverture de la fenêtre deux fois.
            e.stopPropagation();
            fileInput.click();
        });
        
        // Gestion du glisser-déposer (Drag & Drop)
        uploadBox.addEventListener('dragover', (e) => {
            e.preventDefault();
            uploadBox.classList.add('dragover');
        });
        uploadBox.addEventListener('dragleave', () => {
            uploadBox.classList.remove('dragover');
        });
        uploadBox.addEventListener('drop', (e) => {
            e.preventDefault();
            uploadBox.classList.remove('dragover');
            const files = e.dataTransfer.files;
            if (files.length > 0) {
                uploadFile(files[0]);
            }
        });
        
        // Gestion de la sélection de fichier
        fileInput.addEventListener('change', (e) => {
            if (e.target.files.length > 0) {
                uploadFile(e.target.files[0]);
            }
        });
        
        function uploadFile(file) {
            const formData = new FormData();
            formData.append('file', file);
            
            updateStatus('processing', 'Upload en cours...');
            
            fetch('/upload', {
                method: 'POST',
                body: formData
            })
            .then(response => response.json())
            .then(data => {
                if (data.error) {
                    updateStatus('failed', data.error);
                } else {
                    currentTaskId = data.task_id;
                    document.getElementById('extractedText').textContent = data.extracted_text;
                    document.getElementById('logsSection').classList.remove('hidden');
                    updateStatus('processing', 'Résolution en cours...');
                }
            })
            .catch(error => {
                updateStatus('failed', 'Erreur lors de l\'upload');
                console.error('Error:', error);
            });
        }
        
        function updateStatus(status, message) {
            const badge = document.getElementById('statusBadge');
            const info = document.getElementById('taskInfo');
            
            badge.style.display = 'inline-block';
            badge.className = `status-badge status-${status}`;
            badge.textContent = status === 'processing' ? 'En cours' : 
                              status === 'completed' ? 'Terminé' : 'Échec';
            
            info.style.display = 'block';
            info.textContent = message;
        }
        
        function addLog(timestamp, level, message) {
            const container = document.getElementById('logContainer');
            const logEntry = document.createElement('div');
            logEntry.className = `log-entry log-${level}`;
            logEntry.innerHTML = `<span style="color: #666;">[${timestamp}]</span> ${message}`;
            container.appendChild(logEntry);
            container.scrollTop = container.scrollHeight;
        }
        
        // WebSocket events
        socket.on('log_update', (data) => {
            if (data.task_id === currentTaskId) {
                addLog(data.log.timestamp, data.log.level, data.log.message);
            }
        });
        
        socket.on('task_completed', (data) => {
            if (data.task_id === currentTaskId) {
                updateStatus('completed', 'Solution générée avec succès !');
                const downloadSection = document.getElementById('downloadSection');
                downloadSection.style.display = 'block';
                
                document.getElementById('downloadBtn').onclick = () => {
                    window.location.href = `/download/${currentTaskId}`;
                };
            }
        });
        
        socket.on('task_failed', (data) => {
            if (data.task_id === currentTaskId) {
                updateStatus('failed', 'Échec de la résolution (solution partielle disponible)');
                const downloadSection = document.getElementById('downloadSection');
                downloadSection.style.display = 'block';
                
                document.getElementById('downloadBtn').onclick = () => {
                    window.location.href = `/download/${currentTaskId}`;
                };
            }
        });
        
        socket.on('task_error', (data) => {
            if (data.task_id === currentTaskId) {
                updateStatus('failed', `Erreur: ${data.error}`);
            }
        });
    </script>
</body>
</html>
    
    '''
    return template_content

if __name__ == '__main__':
    # Créer le template dans le dossier templates
    os.makedirs('templates', exist_ok=True)
    with open('templates/index.html', 'w', encoding='utf-8') as f:
        # Récupérer le contenu du template depuis la route
        from flask import Flask
        temp_app = Flask(__name__)
        with temp_app.app_context():
            template_content = get_template()
            f.write(template_content)
    
    print("🚀 Démarrage de l'application Flask...")
    print("📝 Template HTML créé dans templates/index.html")
    print("🌐 Application disponible sur http://localhost:5000")
    
    socketio.run(app, debug=True, host='0.0.0.0', port=5000)