image_routes.py

from flask import Blueprint, send_from_directory, current_app, request, jsonify
from flask_login import login_required, current_user
from utils import get_db_connection
import os
import base64
from datetime import datetime

image_bp = Blueprint('image_bp', __name__)

@image_bp.route('/processed/<path:filename>')
def serve_processed_image(filename):
    current_app.logger.info(f"Serving processed image: {filename}")
    return send_from_directory(current_app.config['PROCESSED_FOLDER'], filename)

@image_bp.route('/tmp/<path:filename>')
def serve_tmp_image(filename):
    current_app.logger.info(f"Serving temporary image: {filename}")
    return send_from_directory(current_app.config['TEMP_FOLDER'], filename)

# Proxy routes for /neetprep/processed and /neetprep/tmp
@image_bp.route('/neetprep/processed/<path:filename>')
def serve_neetprep_processed_image(filename):
    current_app.logger.info(f"Serving /neetprep/processed image: {filename}")
    return send_from_directory(current_app.config['PROCESSED_FOLDER'], filename)

@image_bp.route('/neetprep/tmp/<path:filename>')
def serve_neetprep_tmp_image(filename):
    current_app.logger.info(f"Serving /neetprep/tmp image: {filename}")
    return send_from_directory(current_app.config['TEMP_FOLDER'], filename)

@image_bp.route('/save_note_image', methods=['POST'])
@login_required
def save_note_image():
    try:
        if 'image' not in request.files:
            return jsonify({'error': 'No image file provided'}), 400
            
        file = request.files['image']
        image_id = request.form.get('image_id')
        session_id = request.form.get('session_id')
        
        if not image_id or not session_id:
             return jsonify({'error': 'Missing image_id or session_id'}), 400

        # Validate ownership
        conn = get_db_connection()
        img = conn.execute("SELECT i.id, s.user_id FROM images i JOIN sessions s ON i.session_id = s.id WHERE i.id = ?", (image_id,)).fetchone()
        
        if not img or img['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized or image not found'}), 403

        # Save the file
        filename = f"note_{session_id}_{image_id}_{int(datetime.now().timestamp())}.png"
        save_path = os.path.join(current_app.config['PROCESSED_FOLDER'], filename)
        file.save(save_path)
        
        # Update DB
        conn.execute("UPDATE images SET note_filename = ? WHERE id = ?", (filename, image_id))
        conn.commit()
        conn.close()
        
        return jsonify({'success': True, 'filename': filename})

    except Exception as e:
        current_app.logger.error(f"Error saving note: {e}")
        return jsonify({'error': str(e)}), 500

@image_bp.route('/toggle_note_in_pdf', methods=['POST'])
@login_required
def toggle_note_in_pdf():
    """Toggle whether a note should be included in the generated PDF."""
    try:
        data = request.json
        image_id = data.get('image_id')
        include = data.get('include', True)

        if not image_id:
            return jsonify({'error': 'Missing image_id'}), 400

        conn = get_db_connection()
        img = conn.execute("""
            SELECT i.id, s.user_id FROM images i
            JOIN sessions s ON i.session_id = s.id WHERE i.id = ?
        """, (image_id,)).fetchone()

        if not img or img['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403

        conn.execute("UPDATE images SET include_note_in_pdf = ? WHERE id = ?",
                     (1 if include else 0, image_id))
        conn.commit()
        conn.close()

        return jsonify({'success': True, 'include': include})

    except Exception as e:
        current_app.logger.error(f"Error toggling note in PDF: {e}")
        return jsonify({'error': str(e)}), 500

@image_bp.route('/delete_note', methods=['POST'])
@login_required
def delete_note():
    """Delete a note image for a question."""
    try:
        data = request.json
        image_id = data.get('image_id')

        if not image_id:
            return jsonify({'error': 'Missing image_id'}), 400

        conn = get_db_connection()
        img = conn.execute("""
            SELECT i.id, i.note_filename, s.user_id FROM images i
            JOIN sessions s ON i.session_id = s.id WHERE i.id = ?
        """, (image_id,)).fetchone()

        if not img or img['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403

        # Delete the file if it exists
        if img['note_filename']:
            note_path = os.path.join(current_app.config['PROCESSED_FOLDER'], img['note_filename'])
            if os.path.exists(note_path):
                os.remove(note_path)

        conn.execute("UPDATE images SET note_filename = NULL, note_json = NULL WHERE id = ?", (image_id,))
        conn.commit()
        conn.close()

        return jsonify({'success': True})

    except Exception as e:
        current_app.logger.error(f"Error deleting note: {e}")
        return jsonify({'error': str(e)}), 500


@image_bp.route('/save_note_json', methods=['POST'])
@login_required
def save_note_json():
    """Save revision notes as JSON and rasterized PNG for PDF/quiz display."""
    try:
        data = request.json
        image_id = data.get('image_id')
        session_id = data.get('session_id')
        json_data = data.get('json_data')
        image_data = data.get('image_data')

        if not image_id or not session_id or not json_data:
            return jsonify({'error': 'Missing required fields'}), 400

        # Validate ownership
        conn = get_db_connection()
        img = conn.execute("""
            SELECT i.id, s.user_id FROM images i
            JOIN sessions s ON i.session_id = s.id WHERE i.id = ?
        """, (image_id,)).fetchone()

        if not img or img['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403

        # Save JSON to database
        conn.execute("UPDATE images SET note_json = ? WHERE id = ?", (json_data, image_id))

        # Save rasterized PNG if provided (needed for PDF generation and quiz display)
        if image_data and image_data.startswith('data:image/'):
            # Strip the data URL prefix (e.g. "data:image/png;base64,")
            header, encoded = image_data.split(',', 1)
            img_bytes = base64.b64decode(encoded)

            filename = f"note_{session_id}_{image_id}_{int(datetime.now().timestamp())}.png"
            save_path = os.path.join(current_app.config['PROCESSED_FOLDER'], filename)
            with open(save_path, 'wb') as f:
                f.write(img_bytes)

            conn.execute("UPDATE images SET note_filename = ? WHERE id = ?", (filename, image_id))

        conn.commit()
        conn.close()

        return jsonify({'success': True})

    except Exception as e:
        current_app.logger.error(f"Error saving note JSON: {e}")
        return jsonify({'error': str(e)}), 500


@image_bp.route('/get_note_json/<int:image_id>')
@login_required
def get_note_json(image_id):
    """Get revision notes as JSON."""
    try:
        conn = get_db_connection()
        img = conn.execute("""
            SELECT i.note_json, s.user_id FROM images i
            JOIN sessions s ON i.session_id = s.id WHERE i.id = ?
        """, (image_id,)).fetchone()

        if not img or img['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403

        conn.close()

        if img['note_json']:
            return jsonify({'success': True, 'json_data': img['note_json']})
        else:
            return jsonify({'success': False, 'error': 'No note found'}), 404

    except Exception as e:
        current_app.logger.error(f"Error getting note JSON: {e}")
        return jsonify({'error': str(e)}), 500