app.py

from flask import Flask, request, jsonify, render_template, session, redirect, url_for
from flask_session import Session
import google.generativeai as genai
import json
import uuid
import os
from utils.ai_helpers import generate_notebook, stream_notebook_generation, stream_notebook_edit, edit_notebook
from utils.notebook_helpers import format_notebook, extract_notebook_info

app = Flask(__name__)
app.config["SECRET_KEY"] = os.environ.get("SECRET_KEY", "notegenie-secret-key-change-in-production")
app.config["SESSION_TYPE"] = "filesystem"
app.config["SESSION_PERMANENT"] = True
app.config["SESSION_USE_SIGNER"] = True
app.config["PERMANENT_SESSION_LIFETIME"] = 60 * 60 * 24 * 30  # 30 days
app.config["SESSION_FILE_DIR"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), "flask_session")
os.makedirs(app.config["SESSION_FILE_DIR"], exist_ok=True)  # Ensure directory exists

Session(app)

# Map front-end model names to API model names
MODEL_MAPPING = {
    "gemini-2.0-pro": "gemini-2.0-pro-exp-02-05",
    "gemini-2.0-flash": "gemini-2.0-flash",
    "gemini-2.0-flash-thinking": "gemini-2.0-flash-thinking-exp-01-21"
}

def get_api_model_name(frontend_model_name):
    return MODEL_MAPPING.get(frontend_model_name, frontend_model_name)

@app.route("/", methods=["GET"])
def index():
    return render_template("index.html")

@app.route("/set_api_key", methods=["POST"])
def set_api_key():
    api_key = request.form.get("api_key")
    if not api_key:
        return jsonify({"success": False, "message": "API key is required"}), 400
    
    try:
        # Test the API key
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
        response = model.generate_content("Say 'API key is valid'")
        
        # Store API key in session with permanent flag
        session.permanent = True
        session["api_key"] = api_key
        
        return jsonify({"success": True})
    except Exception as e:
        return jsonify({"success": False, "message": str(e)}), 400

@app.route("/generate_notebook", methods=["GET", "POST"])
def generate_notebook_route():
    if "api_key" not in session:
        return jsonify({"success": False, "message": "API key not set"}), 401
    
    # Handle both GET (for streaming) and POST requests
    if request.method == "GET":
        prompt = request.args.get("prompt")
        model_name = request.args.get("model", "gemini-2.0-pro")
        stream = request.args.get("stream", "false").lower() == "true"
    else:
        data = request.json
        prompt = data.get("prompt")
        model_name = data.get("model", "gemini-2.0-pro")
        stream = data.get("stream", False)
        format_only = data.get("format_only", False)
    
    if not prompt:
        return jsonify({"success": False, "message": "Prompt is required"}), 400
    
    # Map the frontend model name to the API model name
    api_model_name = get_api_model_name(model_name)
    
    genai.configure(api_key=session["api_key"])
    
    try:
        # OPTIMIZATION: If format_only is True, skip the AI call and just format the provided content
        if request.method == "POST" and format_only:
            # Use client-provided content as is (it's already the AI response)
            notebook_content = prompt
            notebook_json = format_notebook(notebook_content)
            notebook_info = extract_notebook_info(notebook_content)
            
            return jsonify({
                "success": True,
                "notebook": notebook_json,
                "name": notebook_info["name"],
                "description": notebook_info["description"]
            })
        elif stream:
            return stream_notebook_generation(prompt, api_model_name)
        else:
            notebook_content = generate_notebook(prompt, api_model_name)
            notebook_json = format_notebook(notebook_content)
            notebook_info = extract_notebook_info(notebook_content)
            
            return jsonify({
                "success": True,
                "notebook": notebook_json,
                "name": notebook_info["name"],
                "description": notebook_info["description"]
            })
    except Exception as e:
        return jsonify({"success": False, "message": str(e)}), 500

@app.route("/prepare_edit_notebook", methods=["POST"])
def prepare_edit_notebook():
    """Store the notebook in the session for editing."""
    if "api_key" not in session:
        return jsonify({"success": False, "message": "API key not set"}), 401
    
    data = request.json
    notebook_json = data.get("notebook")
    
    if not notebook_json:
        return jsonify({"success": False, "message": "Notebook content is required"}), 400
    
    # Store the notebook in the session for later access
    session["current_notebook"] = notebook_json
    
    return jsonify({"success": True})

@app.route("/edit_notebook", methods=["GET", "POST"])
def edit_notebook_route():
    if "api_key" not in session:
        return jsonify({"success": False, "message": "API key not set"}), 401
    
    # Get edit prompt and current notebook
    if request.method == "GET":
        edit_prompt = request.args.get("edit_prompt")
        model_name = request.args.get("model", "gemini-2.0-pro")
        stream = request.args.get("stream", "true").lower() == "true"
        # For GET streaming requests, get notebook from session
        notebook_json = session.get("current_notebook")
    else:
        data = request.json
        edit_prompt = data.get("edit_prompt")
        notebook_json = data.get("notebook")
        model_name = data.get("model", "gemini-2.0-pro")
        stream = data.get("stream", False)
    
    if not edit_prompt:
        return jsonify({"success": False, "message": "Edit prompt is required"}), 400
    
    if not notebook_json:
        return jsonify({"success": False, "message": "No notebook available for editing. Please prepare the notebook first."}), 400
    
    # Map the frontend model name to the API model name
    api_model_name = get_api_model_name(model_name)
    
    genai.configure(api_key=session["api_key"])
    
    try:
        if stream:
            return stream_notebook_edit(edit_prompt, notebook_json, api_model_name)
        else:
            # Non-streaming path (not used in current UI but kept for API completeness)
            edited_content = edit_notebook(edit_prompt, notebook_json, api_model_name)
            notebook_json = format_notebook(edited_content)
            notebook_info = extract_notebook_info(edited_content)
            
            return jsonify({
                "success": True,
                "notebook": notebook_json,
                "name": notebook_info["name"],
                "description": notebook_info["description"]
            })
    except Exception as e:
        app.logger.error(f"Error editing notebook: {str(e)}")
        return jsonify({"success": False, "message": str(e)}), 500

@app.route("/download_notebook", methods=["POST"])
def download_notebook():
    from flask import Response
    
    data = request.json
    notebook_json = data.get("notebook")
    filename = data.get("filename", f"notebook_{uuid.uuid4()}.ipynb")
    
    if not notebook_json:
        return jsonify({"success": False, "message": "Notebook content is required"}), 400
    
    if not filename.endswith(".ipynb"):
        filename += ".ipynb"
    
    response = Response(
        json.dumps(notebook_json, indent=2),
        mimetype="application/json",
        headers={"Content-Disposition": f"attachment;filename={filename}"}
    )
    
    return response

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host="0.0.0.0", port=port, debug=(os.environ.get("FLASK_ENV") == "development"))