graph_routes.py

"""
Graph Editor Routes
Links PDF pages to questions as revision notes
"""

from flask import Blueprint, render_template, request, jsonify, current_app, url_for, redirect
from flask_login import login_required, current_user
from utils import get_db_connection, parse_note_payload, dump_note_payload
import os
import json
import fitz
from datetime import datetime
from werkzeug.utils import secure_filename

graph_bp = Blueprint('graph_bp', __name__)

ALLOWED_EXTENSIONS = {'pdf'}

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@graph_bp.route('/session/<session_id>/graph')
@login_required
def graph_editor(session_id):
    """Open graph editor for linking PDF pages to questions."""
    conn = get_db_connection()
    
    session = conn.execute(
        "SELECT * FROM sessions WHERE id = ? AND user_id = ?",
        (session_id, current_user.id)
    ).fetchone()
    
    if not session:
        conn.close()
        return redirect('/dashboard')
    
    # Get question IMAGES (the cropped question images)
    images = conn.execute("""
        SELECT 
            i.id as image_id,
            i.image_index,
            i.filename,
            i.processed_filename,
            i.note_json,
            q.id as question_id,
            q.question_number,
            q.subject,
            q.chapter
        FROM images i
        LEFT JOIN questions q ON i.id = q.image_id
        WHERE i.session_id = ? AND i.image_type = 'cropped'
        ORDER BY i.image_index
    """, (session_id,)).fetchall()
    
    # Get PDF pages (solutions PDF pages stored with image_type='pdf_page')
    pdf_pages = conn.execute("""
        SELECT 
            id as image_id,
            image_index as page_number,
            filename,
            processed_filename
        FROM images
        WHERE session_id = ? AND image_type = 'pdf_page'
        ORDER BY image_index
    """, (session_id,)).fetchall()
    
    # Format PDF pages
    pages_data = []
    for page in pdf_pages:
        page_dict = dict(page)
        # Use processed_filename if available, else original
        if page_dict.get('processed_filename'):
            image_url = url_for('image_bp.serve_processed_image', filename=page_dict['processed_filename'])
        elif page_dict.get('filename'):
            image_url = url_for('main.serve_image', folder='uploads', filename=page_dict['filename'])
        else:
            image_url = f'/placeholder_page/{session_id}/{page_dict["page_number"]+1}'
        
        pages_data.append({
            'image_id': page_dict['image_id'],
            'page_number': page_dict['page_number'] + 1,
            'image_url': image_url,
            'filename': page_dict['filename']
        })
    
    # Format question IMAGES with thumbnails
    images_data = []
    existing_links = []
    for img in images:
        img_dict = dict(img)
        # Use processed_filename for thumbnail (the cropped question image)
        thumb_url = None
        if img_dict.get('processed_filename'):
            thumb_url = url_for('image_bp.serve_processed_image', filename=img_dict['processed_filename'])
        elif img_dict.get('filename'):
            thumb_url = url_for('main.serve_image', folder='uploads', filename=img_dict['filename'])

        note_payload = parse_note_payload(img_dict.get('note_json'))
        linked_pages = note_payload.get('linked_pdf_pages', [])
        
        images_data.append({
            'image_id': img_dict['image_id'],
            'question_number': img_dict.get('question_number'),
            'subject': img_dict.get('subject'),
            'chapter': img_dict.get('chapter'),
            'thumb_url': thumb_url,
            'question_id': img_dict.get('question_id'),
            'linked_page_ids': [page.get('source_page_id') for page in linked_pages if page.get('source_page_id')]
        })

        for linked_page in linked_pages:
            if linked_page.get('source_page_id'):
                existing_links.append({
                    'question_image_id': img_dict['image_id'],
                    'page_id': linked_page.get('source_page_id'),
                    'page_number': linked_page.get('source_page_number')
                })
    
    conn.close()
    
    return render_template('graph_editor.html',
                         session_id=session_id,
                         questions=images_data,  # These are question IMAGES
                         pdf_pages=pages_data,
                         existing_links=existing_links)


@graph_bp.route('/session/graph/upload_pdf', methods=['POST'])
@login_required
def upload_pdf():
    """Upload a solutions PDF and convert pages to images."""
    try:
        if 'pdf' not in request.files:
            return jsonify({'error': 'No PDF file provided'}), 400
        
        file = request.files['pdf']
        session_id = request.form.get('session_id')
        
        if not session_id or file.filename == '':
            return jsonify({'error': 'Missing file or session_id'}), 400
        
        if not allowed_file(file.filename):
            return jsonify({'error': 'Only PDF files allowed'}), 400
        
        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
        
        existing_pages = conn.execute("""
            SELECT id, filename, processed_filename
            FROM images
            WHERE session_id = ? AND image_type = 'pdf_page'
        """, (session_id,)).fetchall()

        for page in existing_pages:
            for stored_name, folder in (
                (page['processed_filename'], current_app.config['PROCESSED_FOLDER']),
                (page['filename'], current_app.config['UPLOAD_FOLDER'])
            ):
                if stored_name:
                    stored_path = os.path.join(folder, stored_name)
                    if os.path.exists(stored_path):
                        try:
                            os.remove(stored_path)
                        except OSError:
                            current_app.logger.warning(f"Could not delete old graph asset: {stored_path}")

        conn.execute("DELETE FROM images WHERE session_id = ? AND image_type = 'pdf_page'", (session_id,))

        # Save uploaded PDF
        pdf_filename = f"solutions_{session_id}_{secure_filename(file.filename)}"
        pdf_path = os.path.join(current_app.config['UPLOAD_FOLDER'], pdf_filename)
        file.save(pdf_path)

        doc = fitz.open(pdf_path)
        page_count = len(doc)
        if page_count == 0:
            doc.close()
            conn.close()
            return jsonify({'error': 'Uploaded PDF has no pages'}), 400

        dpi = getattr(current_user, 'dpi', 150) or 150
        rendered_pages = []

        try:
            for i, page in enumerate(doc):
                pix = page.get_pixmap(dpi=dpi)
                img_filename = f"solutions_{session_id}_page_{i + 1}.png"
                img_path = os.path.join(current_app.config['PROCESSED_FOLDER'], img_filename)
                pix.save(img_path)
                rendered_pages.append(img_path)

                conn.execute('''
                    INSERT INTO images (session_id, image_index, filename, processed_filename, image_type, original_name)
                    VALUES (?, ?, ?, ?, 'pdf_page', ?)
                ''', (session_id, i, pdf_filename, img_filename, f'Solution Page {i + 1}'))
        except Exception:
            for rendered_path in rendered_pages:
                if os.path.exists(rendered_path):
                    try:
                        os.remove(rendered_path)
                    except OSError:
                        current_app.logger.warning(f"Could not clean up rendered graph page: {rendered_path}")
            raise
        finally:
            doc.close()
        
        conn.commit()
        conn.close()
        
        # Redirect back to graph editor (will now have pages)
        return jsonify({
            'success': True,
            'pdf_filename': pdf_filename,
            'redirect': f'/session/{session_id}/graph'
        })
        
    except Exception as e:
        current_app.logger.error(f"Error uploading PDF: {e}")
        return jsonify({'error': str(e)}), 500


@graph_bp.route('/session/graph/save', methods=['POST'])
@login_required
def save_graph():
    """Save links between PDF pages and questions."""
    try:
        data = request.json
        session_id = data.get('session_id')
        links = data.get('links', [])
        
        if not session_id:
            return jsonify({'error': 'Session ID required'}), 400
        
        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
        
        grouped_links = {}
        for link in links:
            page_image_id = link.get('page_id')
            question_image_id = link.get('question_image_id')
            page_number = link.get('page_number')

            if not page_image_id or not question_image_id:
                continue

            page_info = conn.execute("""
                SELECT filename, processed_filename
                FROM images
                WHERE id = ? AND session_id = ?
            """, (page_image_id, session_id)).fetchone()

            if not page_info:
                continue

            grouped_links.setdefault(question_image_id, [])
            grouped_links[question_image_id].append({
                'source_page_id': page_image_id,
                'source_page_number': page_number,
                'source_filename': page_info['filename'],
                'note_filename': page_info['processed_filename'],
                'linked_at': datetime.now().isoformat()
            })

        for question_image_id, linked_pages in grouped_links.items():
            existing = conn.execute("""
                SELECT note_json, note_filename
                FROM images
                WHERE id = ? AND session_id = ?
            """, (question_image_id, session_id)).fetchone()

            existing_payload = parse_note_payload(existing['note_json'] if existing else None)
            existing_by_page = {
                page['source_page_id']: page
                for page in existing_payload['linked_pdf_pages']
                if page.get('source_page_id')
            }

            for page in linked_pages:
                existing_by_page[page['source_page_id']] = page

            merged_pages = sorted(
                existing_by_page.values(),
                key=lambda page: (page.get('source_page_number') or 0, page.get('source_page_id') or 0)
            )
            primary_note_filename = merged_pages[0].get('note_filename') if merged_pages else (existing['note_filename'] if existing else None)
            existing_linked_filenames = {
                page.get('note_filename')
                for page in existing_payload['linked_pdf_pages']
                if page.get('note_filename')
            }
            should_update_note_filename = (
                not existing
                or not existing['note_filename']
                or existing['note_filename'] in existing_linked_filenames
            )

            conn.execute("""
                UPDATE images
                SET note_json = ?,
                    note_filename = ?
                WHERE id = ?
            """, (
                dump_note_payload(
                    existing['note_json'] if existing else None,
                    linked_pdf_pages=merged_pages,
                    annotation_json=existing_payload.get('annotation_json', {})
                ),
                primary_note_filename if should_update_note_filename else existing['note_filename'],
                question_image_id
            ))
        
        conn.commit()
        conn.close()
        
        return jsonify({
            'success': True,
            'message': f'Linked {len(links)} pages to questions',
            'links_count': len(links)
        })
        
    except Exception as e:
        current_app.logger.error(f"Error saving graph: {e}")
        return jsonify({'error': str(e)}), 500


@graph_bp.route('/placeholder_page/<session_id>/<int:page_num>')
@login_required
def placeholder_page(session_id, page_num):
    """Return a placeholder image for PDF pages."""
    from flask import send_file
    from PIL import Image, ImageDraw
    import io
    
    img = Image.new('RGB', (400, 500), color='#1e2130')
    draw = ImageDraw.Draw(img)
    draw.text((200, 250), f"Page {page_num}", fill='#64748b', anchor='mm')
    
    buffer = io.BytesIO()
    img.save(buffer, format='PNG')
    buffer.seek(0)
    
    return send_file(buffer, mimetype='image/png')