from flask import Flask, request, render_template, send_from_directory
import google.generativeai as genai
import os
from PIL import Image
import subprocess
import uuid
import re
import tempfile
from pathlib import Path

app = Flask(__name__)

# Configuration constants
UPLOAD_FOLDER = Path('uploads')
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

# Create uploads directory if it doesn't exist
UPLOAD_FOLDER.mkdir(exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Gemini API configuration
def configure_gemini():
    token = os.environ.get("TOKEN")
    if not token:
        raise ValueError("Environment variable TOKEN must be set.")
    
    genai.configure(api_key=token)
    
    generation_config = {
        "temperature": 1,
        "top_p": 0.95,
        "top_k": 64,
        "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"},
    ]
    
    return genai.GenerativeModel(
        model_name="gemini-1.5-pro",
        generation_config=generation_config,
        safety_settings=safety_settings,
    )

PROMPT_TEMPLATE = """
Résous cet exercice. Tu répondras en détaillant au maximum ton procédé de calcul. 
Réponse attendue uniquement en LaTeX
"""

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

def generate_svg_from_chemfig(chemfig_code, tmpdirname):
    unique_id = str(uuid.uuid4())
    tex_filename = Path(tmpdirname) / f"chem_{unique_id}.tex"
    pdf_filename = Path(tmpdirname) / f"chem_{unique_id}.pdf"
    svg_filename = f"chem_{unique_id}.svg"

    tex_content = f"""\\documentclass[margin=10pt]{{standalone}}
\\usepackage{{chemfig}}
\\begin{{document}}
\\chemfig{{{chemfig_code}}}
\\end{{document}}"""

    tex_filename.write_text(tex_content)

    try:
        subprocess.run(["pdflatex", "-interaction=nonstopmode", str(tex_filename)], 
                      check=True, cwd=tmpdirname)
        subprocess.run(["pdf2svg", str(pdf_filename), 
                      str(UPLOAD_FOLDER / svg_filename)], check=True)
        return svg_filename
    except subprocess.CalledProcessError as e:
        raise Exception(f"LaTeX compilation or SVG conversion error: {e}")

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method != "POST":
        return render_template("index.html", e="")

    if "image" not in request.files:
        return render_template("index.html", e="No image selected.")

    image_file = request.files["image"]
    if not image_file or not allowed_file(image_file.filename):
        return render_template("index.html", e="Invalid file type. Please upload PNG or JPG.")

    try:
        model = configure_gemini()
        
        with tempfile.NamedTemporaryFile(delete=False) as temp_img:
            image_file.save(temp_img.name)
            image = Image.open(temp_img.name)

            response = model.generate_content([PROMPT_TEMPLATE, image])
            latex_code = response.text
            os.unlink(temp_img.name)

            # Handle chemfig diagrams
            match = re.search(r"\\chemfig\{(.*?)\}", latex_code, re.DOTALL)
            if match:
                chemfig_code = match.group(1)
                with tempfile.TemporaryDirectory() as tmpdirname:
                    svg_filename = generate_svg_from_chemfig(chemfig_code, tmpdirname)
                    return render_template("index.html", 
                                        e=f'<img src="/uploads/{svg_filename}" alt="Chemical structure">')

            # Block potentially unsafe content
            if any(keyword in latex_code.lower() for keyword in 
                  ['.pdf', '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.svg']):
                return render_template("index.html", 
                                    e="Sorry, images and PDFs are not yet supported in the response.")

            return render_template("index.html", e=latex_code)

    except Exception as error:
        return render_template("index.html", e=f"Error processing request: {str(error)}")

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

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