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
import logging
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

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - notegenie - %(levelname)s - %(message)s'
)
logger = logging.getLogger()

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

# Set a more permissive file mode for session files to avoid permission issues
app.config["SESSION_FILE_MODE"] = 0o666

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)

# Function to get API key from different sources
def get_api_key():
    # Try to get from session first
    api_key = session.get("api_key")
    
    # Try to get from request header or param (for direct API calls)
    if not api_key:
        api_key = request.headers.get("X-API-Key") or request.args.get("api_key")
    
    return api_key

@app.route("/", methods=["GET"])
def index():
    # Test session functionality
    session["session_test"] = True
    logger.info(f"Session check - variables: {list(session.keys())}")
    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 only
        session.permanent = True
        session["api_key"] = api_key
        
        logger.info("API key successfully set and validated")
        return jsonify({"success": True})
    except Exception as e:
        logger.error(f"API key validation error: {str(e)}")
        return jsonify({"success": False, "message": str(e)}), 400

@app.route("/generate_notebook", methods=["GET", "POST"])
def generate_notebook_route():
    api_key = get_api_key()
    if not api_key:
        logger.warning("Generate notebook request without API key")
        return jsonify({"success": False, "message": "API key not set"}), 401
    
    # Always configure genai with the API key for each request
    genai.configure(api_key=api_key)
    
    # 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)
    
    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:
        logger.error(f"Error generating notebook: {str(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."""
    api_key = get_api_key()
    if not api_key:
        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"] = json.dumps(notebook_json)  # Store as JSON string
    
    return jsonify({"success": True})

@app.route("/edit_notebook", methods=["GET", "POST"])
def edit_notebook_route():
    api_key = get_api_key()
    if not api_key:
        return jsonify({"success": False, "message": "API key not set"}), 401
    
    # Always configure genai with the API key for each request
    genai.configure(api_key=api_key)
    
    # 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")
        if notebook_json:
            notebook_json = json.loads(notebook_json)  # Parse JSON string back to dict
    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)
    
    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

# Add a session diagnostic endpoint
@app.route("/check_session", methods=["GET"])
def check_session():
    # For debugging only - would be disabled in production
    session_data = {
        "has_api_key": "api_key" in session,
        "session_vars": list(session.keys()),
        "session_file_dir_exists": os.path.exists(app.config["SESSION_FILE_DIR"]),
        "session_file_dir_writable": os.access(app.config["SESSION_FILE_DIR"], os.W_OK),
    }
    
    # Check if running on Hugging Face Spaces
    is_hf_space = "SPACE_ID" in os.environ
    session_data["is_huggingface_space"] = is_hf_space
    
    if is_hf_space:
        logger.info("Running on Hugging Face Spaces environment")
    
    return jsonify(session_data)

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