from flask import Flask, request, render_template, jsonify, send_from_directory
from PIL import Image
import google.generativeai as genai
import os
import re
import matplotlib.pyplot as plt
import tempfile
from gradio_client import Client, handle_file
# import subprocess  # Not used
from dataclasses import dataclass
from typing import List, Optional
import logging

# Logging configuration
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class GeminiConfig:
    api_key: str
    generation_config: dict
    safety_settings: List[dict]
    model_name: str = "gemini-exp-1206"

class MathSolver:
    def __init__(self, gemini_config: GeminiConfig):
        self.gemini_config = gemini_config
        genai.configure(api_key=gemini_config.api_key)
        plt.switch_backend('Agg')  # Non-interactive backend

    def query_gemini(self, image_path: str, prompt: str) -> str:
        try:
            img = Image.open(image_path)
            model = genai.GenerativeModel(
                model_name=self.gemini_config.model_name,
                generation_config=self.gemini_config.generation_config,
                safety_settings=self.gemini_config.safety_settings
            )
            response = model.generate_content([prompt, img], request_options={"timeout": 600})
            return response.text
        except Exception as e:
            logger.error(f"Gemini Error: {str(e)}")
            raise

    @staticmethod
    def query_qwen2(image_path: str, question: str) -> str:
        try:
            client = Client("Qwen/Qwen2.5-Math-Demo")
            return client.predict(
                image=handle_file(image_path),
                sketchpad=None,
                question=question,
                api_name="/math_chat_bot"
            )
        except Exception as e:
            logger.error(f"Qwen2 Error: {str(e)}")
            raise

    @staticmethod
    def extract_and_execute_python_code(text: str) -> Optional[List[str]]:
        code_blocks = re.findall(r'```python\n(.*?)```', text, re.DOTALL)
        if not code_blocks:
            return None

        image_paths = []
        for code in code_blocks:
            try:
                code = "import numpy as np\n" + code 
                # Replace single backslashes with double backslashes
                code = code.replace("\\", "\\\\")

                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
                    plt.figure()
                    exec(code)
                    plt.savefig(tmpfile.name)
                    plt.close()
                    relative_path = os.path.basename(tmpfile.name)
                    image_paths.append(relative_path)
            except Exception as e:
                logger.error(f"Error generating graph: {str(e)}")
                continue

        return image_paths if image_paths else None

# Application configuration
app = Flask(__name__)

token = os.environ.get("TOKEN")
gemini_config = GeminiConfig(
    token,  # Replace with your actual API key
    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"},
    ]
)

math_solver = MathSolver(gemini_config)

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

@app.route('/upload', methods=['POST'])
def upload_image():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400

    file = request.files['image']
    if not file.filename:
        return jsonify({'error': 'No file selected'}), 400

    model_choice = request.form.get('model_choice', 'gemini')
    custom_instruction = request.form.get('custom_instruction', '')

    prompt = f"Solve this math problem. Provide a complete solution using LaTeX. {custom_instruction}"

    try:
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            file.save(temp_file.name)

            result = (
                math_solver.query_gemini(temp_file.name, prompt)
                if model_choice == "mariam's"
                else math_solver.query_qwen2(temp_file.name, prompt)
            )

            # Extract and generate graphs
            image_paths = math_solver.extract_and_execute_python_code(result)
            os.unlink(temp_file.name)

            return jsonify({
                'result': result,
                'model': model_choice,
                'image_paths': image_paths,
                'temp_dir': tempfile.gettempdir()
            })
    except Exception as e:
        logger.error(f"Error processing: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/temp/<path:filename>')
def serve_temp_image(filename):
    try:
        return send_from_directory(tempfile.gettempdir(), filename)
    except Exception as e:
        logger.error(f"Error sending image: {str(e)}")
        return jsonify({'error': 'Image not found'}), 404

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