from flask import Flask, request, render_template, jsonify, send_from_directory
from PIL import Image
import google.genai as genai
from google.genai import types
import os
import re
import matplotlib.pyplot as plt
import tempfile
from gradio_client import Client, handle_file
from dataclasses import dataclass
from typing import List, Optional, Tuple
import logging
import time

# 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-2.0-flash-thinking-exp-01-21"

class MathSolver:
    def __init__(self, gemini_config: GeminiConfig):
        self.gemini_config = gemini_config
        self.client = genai.Client(
            api_key=gemini_config.api_key,
            http_options={'api_version': 'v1alpha'}
        )
        plt.switch_backend('Agg')

    def query_gemini(self, image_path: str, prompt: str) -> str:
        try:
            img = Image.open(image_path)
            response = self.client.models.generate_content(
                model=self.gemini_config.model_name,
                config={
                    'thinking_config': {'include_thoughts': True},
                    'generation_config': self.gemini_config.generation_config,
                    'safety_settings': self.gemini_config.safety_settings
                },
                contents=[prompt, img]
            )
            
            result = ""
            for part in response.candidates[0].content.parts:
                if not part.thought:
                    result += part.text + "\n"
            return result.strip()
            
        except Exception as e:
            logger.error(f"Gemini Error: {str(e)}")
            raise

    @staticmethod
    def query_qwen2(image_path: str, question: str, max_retries: int = 3) -> Tuple[str, bool]:
        """
        Query Qwen2 model with retry mechanism
        Returns: (result, success)
        """
        for attempt in range(max_retries):
            try:
                client = Client("Qwen/Qwen2.5-Math-Demo")
                result = client.predict(
                    image=handle_file(image_path),
                    sketchpad=None,
                    question=question,
                    api_name="/math_chat_bot"
                )
                return result, True
            except Exception as e:
                logger.error(f"Qwen2 Error (attempt {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)  # Exponential backoff
                continue
        return "Error: Unable to process with Qwen2 model", False

    @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
                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,
    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 with rendering LaTeX. {custom_instruction}"

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

            if model_choice == "gemini":
                result = math_solver.query_gemini(temp_file.name, prompt)
                success = True
            else:
                result, success = math_solver.query_qwen2(temp_file.name, prompt)
                if not success:
                    # Fallback to Gemini if Qwen2 fails
                    logger.info("Falling back to Gemini model")
                    result = math_solver.query_gemini(temp_file.name, prompt)
                    model_choice = "gemini (fallback)"

            # 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)