knowledge_graph_routes.py

"""
Knowledge Graph Routes
Node-based revision notes system using React Flow
"""

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

knowledge_bp = Blueprint('knowledge_bp', __name__)

# Allowed PDF extensions
ALLOWED_EXTENSIONS = {'pdf'}

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


@knowledge_bp.route('/knowledge_graph')
@login_required
def knowledge_graph_index():
    """Main knowledge graph page."""
    return render_template('knowledge_graph.html')


@knowledge_bp.route('/knowledge_graph/upload_pdf', methods=['POST'])
@login_required
def upload_pdf():
    """
    Upload PDF and create session with page images.
    Returns session_id and page information for creating nodes.
    """
    try:
        if 'pdf' not in request.files:
            return jsonify({'error': 'No PDF file provided'}), 400
        
        file = request.files['pdf']
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        if not allowed_file(file.filename):
            return jsonify({'error': 'Only PDF files are allowed'}), 400
        
        # Create session for this knowledge graph
        conn = get_db_connection()
        session_id = f"kg_{current_user.id}_{int(datetime.now().timestamp())}"
        
        # Save session metadata
        conn.execute("""
            INSERT INTO sessions (id, user_id, original_filename, session_type, created_at)
            VALUES (?, ?, ?, 'knowledge_graph', ?)
        """, (session_id, current_user.id, secure_filename(file.filename), datetime.now()))
        
        # Save PDF file
        pdf_filename = f"{session_id}_{secure_filename(file.filename)}"
        pdf_path = os.path.join(current_app.config['UPLOAD_FOLDER'], pdf_filename)
        file.save(pdf_path)
        
        conn.commit()
        
        # TODO: Process PDF to extract pages
        # For now, we'll create placeholder page nodes
        # In production, use PyMuPDF or pdf2image to convert PDF pages to images
        
        pages = []
        # Placeholder: Create 5 dummy pages
        # Replace this with actual PDF processing
        for i in range(1, 6):
            pages.append({
                'page_number': i,
                'image_url': f'/placeholder_page/{session_id}/{i}'  # Placeholder route
            })
        
        conn.close()
        
        return jsonify({
            'success': True,
            'session_id': session_id,
            'pdf_filename': pdf_filename,
            'pages': pages
        })
        
    except Exception as e:
        current_app.logger.error(f"Error uploading PDF: {e}")
        return jsonify({'error': str(e)}), 500


@knowledge_bp.route('/knowledge_graph/save', methods=['POST'])
@login_required
def save_graph():
    """
    Save the node graph (nodes and edges) to database.
    """
    try:
        data = request.json
        session_id = data.get('session_id')
        nodes = data.get('nodes', [])
        edges = data.get('edges', [])
        
        if not session_id:
            return jsonify({'error': 'Session ID required'}), 400
        
        # Validate ownership
        conn = get_db_connection()
        session = conn.execute(
            "SELECT user_id FROM sessions WHERE id = ? AND session_type = 'knowledge_graph'",
            (session_id,)
        ).fetchone()
        
        if not session or session['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403
        
        # Save graph data as JSON
        graph_data = {
            'nodes': nodes,
            'edges': edges,
            'saved_at': datetime.now().isoformat()
        }
        
        conn.execute("""
            INSERT OR REPLACE INTO knowledge_graphs (session_id, graph_data, updated_at)
            VALUES (?, ?, ?)
        """, (session_id, json.dumps(graph_data), datetime.now()))
        
        conn.commit()
        conn.close()
        
        return jsonify({
            'success': True,
            'message': 'Graph saved successfully'
        })
        
    except Exception as e:
        current_app.logger.error(f"Error saving graph: {e}")
        return jsonify({'error': str(e)}), 500


@knowledge_bp.route('/knowledge_graph/load/<session_id>')
@login_required
def load_graph(session_id):
    """
    Load a saved node graph.
    """
    try:
        conn = get_db_connection()
        
        # Validate ownership
        session = conn.execute(
            "SELECT user_id FROM sessions WHERE id = ? AND session_type = 'knowledge_graph'",
            (session_id,)
        ).fetchone()
        
        if not session or session['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403
        
        # Load graph data
        graph = conn.execute(
            "SELECT graph_data FROM knowledge_graphs WHERE session_id = ?",
            (session_id,)
        ).fetchone()
        
        conn.close()
        
        if not graph or not graph['graph_data']:
            return jsonify({
                'success': True,
                'nodes': [],
                'edges': []
            })
        
        graph_data = json.loads(graph['graph_data'])
        
        return jsonify({
            'success': True,
            'nodes': graph_data.get('nodes', []),
            'edges': graph_data.get('edges', []),
            'saved_at': graph_data.get('saved_at')
        })
        
    except Exception as e:
        current_app.logger.error(f"Error loading graph: {e}")
        return jsonify({'error': str(e)}), 500


@knowledge_bp.route('/knowledge_graph/list')
@login_required
def list_graphs():
    """
    List all knowledge graphs for the current user.
    """
    try:
        conn = get_db_connection()
        
        graphs = conn.execute("""
            SELECT s.id, s.original_filename, s.created_at, kg.updated_at
            FROM sessions s
            LEFT JOIN knowledge_graphs kg ON s.id = kg.session_id
            WHERE s.user_id = ? AND s.session_type = 'knowledge_graph'
            ORDER BY s.created_at DESC
        """, (current_user.id,)).fetchall()
        
        conn.close()
        
        return jsonify({
            'success': True,
            'graphs': [dict(g) for g in graphs]
        })
        
    except Exception as e:
        current_app.logger.error(f"Error listing graphs: {e}")
        return jsonify({'error': str(e)}), 500


@knowledge_bp.route('/knowledge_graph/delete/<session_id>', methods=['POST'])
@login_required
def delete_graph(session_id):
    """
    Delete a knowledge graph.
    """
    try:
        conn = get_db_connection()
        
        # Validate ownership
        session = conn.execute(
            "SELECT user_id FROM sessions WHERE id = ? AND session_type = 'knowledge_graph'",
            (session_id,)
        ).fetchone()
        
        if not session or session['user_id'] != current_user.id:
            conn.close()
            return jsonify({'error': 'Unauthorized'}), 403
        
        # Delete graph data and session
        conn.execute("DELETE FROM knowledge_graphs WHERE session_id = ?", (session_id,))
        conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
        
        conn.commit()
        conn.close()
        
        return jsonify({
            'success': True,
            'message': 'Graph deleted successfully'
        })
        
    except Exception as e:
        current_app.logger.error(f"Error deleting graph: {e}")
        return jsonify({'error': str(e)}), 500


@knowledge_bp.route('/placeholder_page/<session_id>/<int:page_num>')
@login_required
def placeholder_page(session_id, page_num):
    """
    Placeholder route for page images.
    Returns a simple placeholder image.
    """
    from flask import send_file
    from PIL import Image, ImageDraw, ImageFont
    import io
    
    # Create a placeholder image
    img = Image.new('RGB', (400, 500), color='#2b3035')
    draw = ImageDraw.Draw(img)
    
    # Draw text
    text = f"Page {page_num}"
    draw.text((200, 250), text, fill='#ffffff', anchor='mm')
    
    # Save to buffer
    buffer = io.BytesIO()
    img.save(buffer, format='PNG')
    buffer.seek(0)
    
    return send_file(buffer, mimetype='image/png')


@knowledge_bp.route('/knowledge_graph/add_question', methods=['POST'])
@login_required
def add_question():
    """
    Add a question node from existing session data.
    """
    try:
        data = request.json
        session_id = data.get('session_id')
        question_data = data.get('question', {})
        
        # Validate and fetch question from database
        conn = get_db_connection()
        question = conn.execute("""
            SELECT q.*, i.filename as image_filename
            FROM questions q
            LEFT JOIN images i ON q.image_id = i.id
            WHERE q.session_id = ? AND q.id = ?
        """, (session_id, question_data.get('id'))).fetchone()
        
        if not question:
            conn.close()
            return jsonify({'error': 'Question not found'}), 404
        
        conn.close()
        
        # Return question data formatted for node creation
        return jsonify({
            'success': True,
            'node': {
                'id': f"question-{question['id']}",
                'type': 'question',
                'data': {
                    'number': question.get('question_number', '?'),
                    'subject': question.get('subject', ''),
                    'chapter': question.get('chapter', ''),
                    'question_id': question['id']
                }
            }
        })
        
    except Exception as e:
        current_app.logger.error(f"Error adding question: {e}")
        return jsonify({'error': str(e)}), 500


@knowledge_bp.route('/knowledge_graph/add_note', methods=['POST'])
@login_required
def add_note():
    """
    Add a note node from existing revision notes.
    """
    try:
        data = request.json
        image_id = data.get('image_id')
        
        conn = get_db_connection()
        
        # Fetch note from database
        image = conn.execute("""
            SELECT note_json, note_filename
            FROM images
            WHERE id = ? AND note_json IS NOT NULL
        """, (image_id,)).fetchone()
        
        if not image:
            conn.close()
            return jsonify({'error': 'Note not found'}), 404
        
        conn.close()
        
        note_content = image['note_json'] or '{}'
        
        return jsonify({
            'success': True,
            'node': {
                'id': f"note-{image_id}",
                'type': 'note',
                'data': {
                    'content': note_content,
                    'image_id': image_id,
                    'note_image': image['note_filename']
                }
            }
        })
        
    except Exception as e:
        current_app.logger.error(f"Error adding note: {e}")
        return jsonify({'error': str(e)}), 500