import os
import sys
import gradio as gr
from PIL import Image
import tempfile
import json
from flask import Flask, request, jsonify
from ctransformers import AutoModelForCausalLM

# Import du module OCR
from ocr_module import process_image as ocr_process

# Chemins
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(CURRENT_DIR, "models", "Mistral-7B-Instruct-v0.3.Q4_K_M.gguf")
TEMP_DIR = os.path.join(CURRENT_DIR, "temp")

# Créer le dossier temp s'il n'existe pas
if not os.path.exists(TEMP_DIR):
    os.makedirs(TEMP_DIR)

# Vérifier l'existence du modèle
if not os.path.exists(MODEL_PATH):
    print(f"ATTENTION: Le modèle n'existe pas à {MODEL_PATH}")
    print("Exécutez d'abord le script download_model.py pour copier le modèle")
    # Tenter de trouver le modèle à l'emplacement alternatif (pour le développement)
    alt_path = os.path.join(os.path.dirname(CURRENT_DIR), "models", "Mistral-7B-Instruct-v0.3.Q4_K_M.gguf")
    if os.path.exists(alt_path):
        print(f"Modèle trouvé à l'emplacement alternatif: {alt_path}")
        MODEL_PATH = alt_path
    else:
        print("Aucun modèle trouvé. L'application risque de ne pas fonctionner correctement.")

# Chargement du modèle Mistral
try:
    print(f"Chargement du modèle depuis: {MODEL_PATH}")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        model_type="mistral",
        local_files_only=True,
        gpu_layers=1
    )
    print("Modèle chargé avec succès")
except Exception as e:
    print(f"Erreur lors du chargement du modèle: {e}")
    model = None

# Fonction pour prétraiter et OCR une image
def process_image(image):
    try:
        # Sauvegarder l'image temporairement
        with tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir=TEMP_DIR) as temp_file:
            temp_path = temp_file.name
            if isinstance(image, str):  # Si c'est un chemin
                Image.open(image).save(temp_path)
            else:  # Si c'est un objet image
                image.save(temp_path)
        
        # Appeler la fonction OCR du module
        ocr_result = ocr_process(temp_path)
        
        # Supprimer le fichier temporaire
        if os.path.exists(temp_path):
            os.remove(temp_path)
        
        return ocr_result
    except Exception as e:
        return {"error": str(e)}

# Fonction pour moderniser le texte avec Mistral
def modernize_text(text):
    try:
        if model is None:
            return {
                "status": "error",
                "error": "Modèle non disponible"
            }
            
        # Créer le prompt pour le modèle
        prompt = f"""<s>[INST] Tu es un assistant spécialisé dans la modernisation de textes anciens. 
Modernise le texte suivant en français contemporain tout en préservant son sens:

{text} [/INST]"""
        
        # Générer la réponse
        response = model(prompt, max_new_tokens=1024, temperature=0.7)
        
        return {
            "status": "success",
            "modern_text": response.strip()
        }
    except Exception as e:
        return {
            "status": "error",
            "error": str(e)
        }

# Fonction principale pour l'interface Gradio
def ocr_and_modernize(image):
    # Extraire le texte de l'image
    ocr_result = process_image(image)
    
    if "text" in ocr_result and ocr_result["text"]:
        # Moderniser le texte
        modernization_result = modernize_text(ocr_result["text"])
        
        if modernization_result["status"] == "success":
            result = {
                "Texte original": ocr_result["text"],
                "Texte modernisé": modernization_result["modern_text"]
            }
            
            # Si LaTeX est disponible
            if "latex" in ocr_result and ocr_result["latex"]:
                result["Formules LaTeX"] = ocr_result["latex"]
                
            # Ajouter la méthode OCR utilisée
            if "ocr_method" in ocr_result:
                result["Méthode OCR"] = ocr_result["ocr_method"]
                
            # Ajouter la confiance OCR
            if "confidence" in ocr_result:
                result["Confiance OCR"] = f"{ocr_result['confidence']:.2f}%"
                
            return result
        else:
            return {
                "Texte original": ocr_result["text"],
                "Erreur": f"Échec de la modernisation: {modernization_result.get('error', 'Erreur inconnue')}"
            }
    else:
        return {
            "Erreur": f"Échec de l'OCR: {ocr_result.get('error', 'Aucun texte détecté')}"
        }

# Interface Gradio
with gr.Blocks(title="OCR et Modernisation de Texte") as demo:
    gr.Markdown("# OCR et Modernisation de Texte")
    gr.Markdown("Téléchargez une image contenant du texte pour extraire et moderniser son contenu.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Image à analyser")
            submit_btn = gr.Button("Analyser l'image", variant="primary")
        
        with gr.Column():
            output = gr.JSON(label="Résultats")
            
    submit_btn.click(fn=ocr_and_modernize, inputs=input_image, outputs=output)
    
    gr.Examples(
        examples=["exemple1.jpg", "exemple2.jpg"] if os.path.exists("exemple1.jpg") else [],
        inputs=input_image
    )

# API Flask pour intégration
app = Flask(__name__)

@app.route("/api/ocr", methods=["POST"])
def api_ocr():
    if "image" not in request.files:
        return jsonify({"error": "Aucune image n'a été fournie"}), 400
    
    file = request.files["image"]
    temp_path = os.path.join(TEMP_DIR, "temp_" + str(hash(file.filename)) + ".png")
    file.save(temp_path)
    
    try:
        # Extraire le texte
        ocr_result = process_image(temp_path)
        
        # Moderniser le texte
        if "text" in ocr_result and ocr_result["text"]:
            modernization_result = modernize_text(ocr_result["text"])
            if modernization_result["status"] == "success":
                ocr_result["original_text"] = ocr_result["text"]
                ocr_result["text"] = modernization_result["modern_text"]
                
        return jsonify(ocr_result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        # Nettoyer
        if os.path.exists(temp_path):
            os.remove(temp_path)

# Point d'entrée
if __name__ == "__main__":
    # Si exécuté sur HF Spaces, utiliser Gradio
    if os.environ.get("SPACE_ID"):
        demo.launch()
    else:
        # En local, démarrer le serveur Flask
        app.run(host="0.0.0.0", port=7860, debug=True)