routes/serving.py

import os
import uuid
import base64
import numpy as np
import cv2
import fitz
from flask import current_app, url_for, send_from_directory, send_file, render_template
from .common import main_bp, get_db_connection, login_required, current_user

@main_bp.route('/tmp/<path:filename>')
@login_required
def serve_tmp_file(filename):
    return send_from_directory(current_app.config['TEMP_FOLDER'], filename)

@main_bp.route('/processed/<path:filename>')
@login_required
def serve_processed_file(filename):
    conn = get_db_connection()
    # Check both processed_filename and note_filename columns
    image_owner = conn.execute(
        "SELECT s.user_id FROM images i JOIN sessions s ON i.session_id = s.id WHERE i.processed_filename = ? OR i.note_filename = ?",
        (filename, filename)
    ).fetchone()
    conn.close()

    if image_owner and image_owner['user_id'] == current_user.id:
        return send_from_directory(current_app.config['PROCESSED_FOLDER'], filename)
    else:
        return "Unauthorized", 403

@main_bp.route('/image/<folder>/<path:filename>')
def serve_image(folder, filename):
    folder_map = {
        'uploads': 'UPLOAD_FOLDER',
        'processed': 'PROCESSED_FOLDER',
        'output': 'OUTPUT_FOLDER'
    }
    config_key = folder_map.get(folder, f'{folder.upper()}_FOLDER')
    base_folder_path = current_app.config.get(config_key)
    
    if not base_folder_path:
        return "Not found", 404

    full_path = os.path.abspath(os.path.join(base_folder_path, filename))
    if not full_path.startswith(os.path.abspath(base_folder_path)):
        return "Forbidden", 403

    if os.path.exists(full_path):
        return send_from_directory(base_folder_path, filename)
    else:
        parts = filename.split('/')
        if len(parts) > 1:
            fallback_filename = parts[-1]
            fallback_full_path = os.path.abspath(os.path.join(base_folder_path, fallback_filename))
            if os.path.exists(fallback_full_path):
                 return send_from_directory(base_folder_path, fallback_filename)

    return "Not found", 404

@main_bp.route('/download/<filename>')
def download_file(filename):
    return send_file(os.path.join(current_app.config['OUTPUT_FOLDER'], filename), as_attachment=True)

@main_bp.route('/view_pdf/<filename>')
def view_pdf(filename):
    return send_file(os.path.join(current_app.config['OUTPUT_FOLDER'], filename), as_attachment=False)

@main_bp.route('/view_pdf_v2/<filename>')
def view_pdf_v2(filename):
    return render_template('pdfjs_viewer.html', pdf_url=url_for('main.view_pdf', filename=filename), pdf_title=filename)

@main_bp.route('/view_generated_pdf/<filename>')
def view_generated_pdf(filename):
    from werkzeug.utils import secure_filename
    safe_filename = secure_filename(filename)
    filepath = os.path.join(current_app.config['TEMP_FOLDER'], safe_filename)
    if not os.path.exists(filepath):
        return "File not found.", 404
    return send_file(filepath, mimetype='application/pdf', as_attachment=False)

@main_bp.route('/viewpdflegacy/<int:pdf_id>')
@login_required
def view_pdf_legacy(pdf_id):
    conn = get_db_connection()
    pdf_info = conn.execute('SELECT filename, subject FROM generated_pdfs WHERE id = ?', (pdf_id,)).fetchone()
    conn.close()

    if not pdf_info:
        return "PDF not found", 404

    pdf_filename = pdf_info['filename']
    pdf_subject = pdf_info['subject']
    pdf_path = os.path.join(current_app.config['OUTPUT_FOLDER'], pdf_filename)

    if not os.path.exists(pdf_path):
        return "PDF file not found on disk", 404

    image_paths = []
    try:
        doc = fitz.open(pdf_path)
        for i in range(0, doc.page_count, 2):
            page1 = doc.load_page(i)
            pix1 = page1.get_pixmap(dpi=current_user.dpi)

            if i + 1 < doc.page_count:
                page2 = doc.load_page(i + 1)
                pix2 = page2.get_pixmap(dpi=current_user.dpi)
                img1 = np.frombuffer(pix1.samples, dtype=np.uint8).reshape(pix1.h, pix1.w, pix1.n)
                img2 = np.frombuffer(pix2.samples, dtype=np.uint8).reshape(pix2.h, pix2.w, pix2.n)
                if img1.shape[2] == 4: img1 = cv2.cvtColor(img1, cv2.COLOR_RGBA2RGB)
                if img2.shape[2] == 4: img2 = cv2.cvtColor(img2, cv2.COLOR_RGBA2RGB)
                max_h = max(img1.shape[0], img2.shape[0])
                if img1.shape[0] < max_h:
                    img1 = np.pad(img1, ((0, max_h - img1.shape[0]), (0, 0), (0, 0)), mode='constant', constant_values=255)
                if img2.shape[0] < max_h:
                    img2 = np.pad(img2, ((0, max_h - img2.shape[0]), (0, 0), (0, 0)), mode='constant', constant_values=255)
                combined_img = np.hstack((img1, img2))
                combined_pix = fitz.Pixmap(fitz.csRGB, combined_img.shape[1], combined_img.shape[0], combined_img.tobytes())
            else:
                combined_pix = pix1

            temp_image_filename = f"legacy_view_{uuid.uuid4()}_page_{i}_{i+1}.png"
            temp_image_path = os.path.join(current_app.config['PROCESSED_FOLDER'], temp_image_filename)
            combined_pix.save(temp_image_path)
            image_paths.append(url_for('main.serve_image', folder='processed', filename=temp_image_filename))
        doc.close()
    except Exception as e:
        return f"Error processing PDF: {str(e)}", 500

    return render_template('simple_viewer.html', image_urls=image_paths, pdf_title=pdf_subject or pdf_filename)