image_routes.py

from flask import Blueprint, send_from_directory, current_app, request, jsonify, url_for
from flask_login import login_required, current_user
from utils import get_db_connection, parse_note_payload, dump_note_payload
import os
import base64
import json
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('/upload_note_reference', methods=['POST'])
@login_required
def upload_note_reference():
    """Upload reference images for revision notes."""
    try:
        if 'image' not in request.files:
            return jsonify({'error': 'No image file provided'}), 400

        file = request.files['image']
        session_id = request.form.get('session_id')
        image_id = request.form.get('image_id')

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

        # Validate ownership
        conn = get_db_connection()
        session = conn.execute("SELECT user_id FROM sessions WHERE id = ?", (session_id,)).fetchone()
        
        if not session or session['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403

        # Save uploaded image
        filename = f"ref_{session_id}_{image_id}_{int(datetime.now().timestamp())}_{file.filename}"
        save_path = os.path.join(current_app.config['TEMP_FOLDER'], filename)
        file.save(save_path)

        conn.close()

        return jsonify({
            'success': True,
            'filename': filename,
            'url': url_for('image_bp.serve_tmp_image', filename=filename)
        })

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

@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, 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

        # Delete the file if it exists
        linked_note_filenames = {
            page.get('note_filename')
            for page in parse_note_payload(img['note_json']).get('linked_pdf_pages', [])
            if page.get('note_filename')
        }

        if img['note_filename'] and img['note_filename'] not in linked_note_filenames:
            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

        existing_row = conn.execute("SELECT note_json FROM images WHERE id = ?", (image_id,)).fetchone()
        existing_payload = parse_note_payload(existing_row['note_json'] if existing_row else None)

        annotation_json = {}
        if isinstance(json_data, str):
            try:
                parsed_json_data = json.loads(json_data)
                if isinstance(parsed_json_data, dict):
                    annotation_json = parsed_json_data
            except (TypeError, ValueError):
                annotation_json = {}
        elif isinstance(json_data, dict):
            annotation_json = json_data

        conn.execute(
            "UPDATE images SET note_json = ? WHERE id = ?",
            (
                dump_note_payload(
                    existing_row['note_json'] if existing_row else None,
                    linked_pdf_pages=existing_payload.get('linked_pdf_pages', []),
                    annotation_json=annotation_json
                ),
                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 and image URL."""
    try:
        conn = get_db_connection()
        img = conn.execute("""
            SELECT i.note_json, 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

        conn.close()

        image_data = None
        if img['note_filename']:
            image_data = url_for('image_bp.serve_processed_image', filename=img['note_filename'])

        parsed_payload = parse_note_payload(img['note_json'])

        if img['note_json'] or image_data:
            return jsonify({
                'success': True, 
                'json_data': json.dumps(parsed_payload.get('annotation_json', {})),
                'image_data': image_data,
                'linked_pages': parsed_payload.get('linked_pdf_pages', [])
            })
        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