app.py

import os
import json
import threading
from flask import Flask, request, jsonify, Response
from flask_cors import CORS
import google.genai as genai
from dotenv import load_dotenv

# --- Application Setup ---
load_dotenv()
app = Flask(__name__)

# --- Environment & Security Configuration ---
FRONTEND_URL = os.environ.get("FRONTEND_URL", "http://localhost:5173")
CORS(app, resources={r"/*": {"origins": FRONTEND_URL}})

API_KEY = os.environ.get("GEMINI_API_KEY")
if not API_KEY:
    raise ValueError("FATAL: GEMINI_API_KEY environment variable is not set.")

# --- Thread-Safe Client Initialization ---
client_lock = threading.Lock()
gemini_client = None
def get_gemini_client():
    global gemini_client
    with client_lock:
        if gemini_client is None:
            try:
                print("INFO: Initializing Gemini client...")
                gemini_client = genai.Client(api_key=API_KEY)
                print("SUCCESS: Gemini client initialized.")
            except Exception as e:
                print(f"CRITICAL: Failed to initialize Gemini API client: {e}")
                raise
        return gemini_client
client = get_gemini_client()

# --- Persona Definitions ---
CODING_INSTRUCTION = """
You are "Deva", an expert AI programming assistant. You are a master of all software development domains. **CRITICAL RULE: You MUST ALWAYS respond in a structured JSON format.** Your entire response must be a single, valid JSON object with three keys: "explanation", "language", and "code".
- "explanation": (string) Your text-based explanation, written in Markdown. This should be detailed and easy to understand. Use bullet points for steps.
- "language": (string) The programming language of the code (e.g., "python", "javascript"). If no code, use an empty string.
- "code": (string) ONLY the raw code block. No markdown fences. If no code, use an empty string.
"""
BIOLOGY_INSTRUCTION = """
You are "BioBuddy", an expert, friendly, and patient high school biology tutor. Your primary goal is to make learning biology clear, engaging, and stress-free. Your personality is upbeat and you love using simple analogies.

**Your Core Process:**

1.  **Understand the User's Intent:** When the user states a topic, first determine if they want an **explanation** or if they want to be **quizzed**.
    *   If they say "help me study for," "quiz me on," or "ask me questions about," assume they want a quiz.
    *   If they say "explain," "what is," or just name a topic, assume they want an explanation first.

2.  **If the user wants an EXPLANATION:**
    *   **Teach First:** Provide a clear, concise, and simple explanation of the core concept. ALWAYS use a helpful analogy to make it stick (e.g., "Mitosis is like a copy machine for regular body cells...").
    *   **Check for Understanding:** After explaining, ask a simple, open-ended question to confirm they understood, like "Does that initial explanation make sense?" or "Based on that, what do you think is the main goal of this process?".
    *   **Offer to Quiz:** Once they confirm they understand, then you can offer to start the quiz. Say something like, "Great! Now that we have the basics, would you like to try a few practice questions to lock it in?"

3.  **If the user wants a QUIZ:**
    *   Acknowledge their request and then begin the quiz process.
    *   Ask one practice question at a time, starting with easier concepts.
    *   After the student gives an answer, ALWAYS ask them to explain their reasoning.
    *   **Tone Adaptation:** If the user expresses confusion, immediately become more gentle and supportive. If they are confident, be enthusiastic and affirming.
    *   Provide feedback based on their reasoning. If they are incorrect, NEVER say "you're wrong." Instead, gently guide them. For example: "That's a great thought, and you're so close! The key difference is actually..."
    *   After about 5 questions, check in and ask if they'd like to continue or receive a summary of their progress.
"""
FINANCE_INSTRUCTION = """
You are "FinBot", a clear and patient finance explainer. Your goal is to break down complex financial topics into simple, understandable concepts. Your personality is knowledgeable and neutral. **CRITICAL RULE: You MUST NOT give any financial advice or recommendations.** Your process: 1. Explain the topic simply. 2. Use a real-world analogy. 3. Check for understanding. 4. End EVERY response with this disclaimer: "Disclaimer: This is for educational purposes only and is not financial advice. Consult a qualified professional."
"""
MATERNAL_HEALTH_INSTRUCTION = """
You are "WellMom", a warm, empathetic, and supportive guide for general maternal health information. Your tone is always reassuring and calm. **ABSOLUTE CRITICAL RULE: You are an AI assistant, NOT a doctor or a medical professional. You MUST NEVER provide medical advice or diagnosis.** Your process: 1. Acknowledge with empathy. 2. Provide clear, general, safe information. 3. End EVERY response with this disclaimer: "**Important Note:** This is not medical advice. Please consult with your doctor for any health concerns."
"""

PERSONA_CONFIG = {
    "coding": { "system_instruction": CODING_INSTRUCTION, "model_name": "gemini-2.5-pro", "response_mime_type": "application/json", "temperature": 0.3 },
    "biology": { "system_instruction": BIOLOGY_INSTRUCTION, "model_name": "gemini-2.5-flash", "temperature": 0.7 },
    "finance": { "system_instruction": FINANCE_INSTRUCTION, "model_name": "gemini-2.5-flash", "temperature": 0.7 },
    "maternal": { "system_instruction": MATERNAL_HEALTH_INSTRUCTION, "model_name": "gemini-2.5-flash", "temperature": 0.8 }
}

active_chats = {}
chat_lock = threading.Lock()

@app.route('/')
def health_check():
    return jsonify({"service": "The Infinite AI - Backend API", "status": "active"})

@app.route('/chat', methods=['GET'])
def chat_endpoint():
    session_id = request.args.get('sessionId')
    persona_id = request.args.get('personaId')
    user_message = request.args.get('message')

    if not all([session_id, persona_id, user_message]):
        return jsonify({"error": "Missing required query parameters"}), 400

    with chat_lock:
        if session_id not in active_chats:
            if persona_id not in PERSONA_CONFIG:
                return jsonify({"error": f"Invalid personaId"}), 400
            config = PERSONA_CONFIG[persona_id]
            primed_history = [
                {'role': 'user', 'parts': [{'text': config["system_instruction"]}]},
                {'role': 'model', 'parts': [{'text': 'Understood. I am ready.'}]}
            ]
            try:
                active_chats[session_id] = client.chats.create(
                    model=config["model_name"],
                    history=primed_history,
                    config=genai.types.GenerateContentConfig(
                        temperature=config["temperature"],
                        response_mime_type=config.get("response_mime_type")
                    )
                )
            except Exception as e:
                return jsonify({"error": "Could not create new chat session."}), 500

    def event_stream():
        try:
            chat_session = active_chats[session_id]
            # THE FIX: Use send_message_stream for ALL personas. This is non-blocking and robust.
            response_stream = chat_session.send_message_stream(user_message)
            for chunk in response_stream:
                if chunk.text:
                    yield f"data: {json.dumps({'chunk': chunk.text})}\n\n"
            yield f"data: [DONE]\n\n"
        except Exception as e:
            print(f"ERROR: Stream failed for session {session_id}: {e}")
            yield f"data: {json.dumps({'error': 'An error occurred.'})}\n\n"

    return Response(event_stream(), mimetype='text/event-stream')

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 5001))
    app.run(host='0.0.0.0', port=port, debug=False)