import os
import requests

API_KEY = os.getenv("API_KEY_GEMINI")
GEMINI_URL = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={API_KEY}"

def filter_math_query(query):
    """Checks if the query is related to a mathematical topic and corrects spelling mistakes using a single API call."""

    if not API_KEY:
        return "⚠️ Error: Missing API key. Please configure API_KEY_GEMINI.", None

    headers = {"Content-Type": "application/json"}

    prompt = (
        f"Analyze the following query and perform two tasks:\n"
        f"1️⃣ Check whether it strictly relates to mathematics.\n"
        f"   - Reply only with 'MATH' if it's a valid math-related topic.\n"
        f"   - Reply with 'NON MATH' if it's not a mathematical subject.\n"
        f"2️⃣ If the query is mathematical, correct any spelling mistakes.\n"
        f"   - Respond in the format 'MATH\nCorrection: [Corrected query]'.\n"
        f"   - If the query is not mathematical, do not provide a correction.\n\n"
        f"Query: {query}"
    )

    data = {
        "contents": [
            {
                "parts": [
                    {"text": prompt}
                ]
            }
        ]
    }

    try:
        response = requests.post(GEMINI_URL, headers=headers, json=data)
        response_json = response.json()
        result = response_json["candidates"][0]["content"]["parts"][0]["text"].strip()

        if result.startswith("MATH"):
            corrected_query = query  # Default to the original query
            if "Correction:" in result:
                corrected_query = result.split("Correction:")[1].strip()

            return "✅ Query accepted as mathematical.", corrected_query

        elif "NON MATH" in result:
            return "❌ Query rejected. This is not a mathematical topic.", None

        else:
            return f"⚠️ Error: Unexpected response from Gemini ({result})", None

    except Exception as e:
        return f"⚠️ Gemini API Error: {e}", None