routes.py

import os
import uuid
import logging
from datetime import datetime
from flask import render_template, request, redirect, url_for, flash, session, jsonify, send_file
from werkzeug.utils import secure_filename
from werkzeug.security import generate_password_hash, check_password_hash
from app import app, db
from models import User, ReviewSession
from edsl_service import EDSLService
import tempfile
import zipfile

# Configure logging
logger = logging.getLogger(__name__)

ALLOWED_EXTENSIONS = {'pdf'}

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

@app.route('/')
def index():
    if 'user_id' not in session:
        return render_template('index.html')
    
    user = User.query.get(session['user_id'])
    recent_sessions = ReviewSession.query.filter_by(user_id=user.id).order_by(ReviewSession.created_at.desc()).limit(5).all()
    return render_template('index.html', user=user, recent_sessions=recent_sessions)

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        user = User.query.filter_by(username=username).first()
        
        if user and check_password_hash(user.password_hash, password):
            session['user_id'] = user.id
            flash('Successfully logged in!', 'success')
            return redirect(url_for('index'))
        else:
            flash('Invalid username or password', 'error')
    
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        email = request.form['email']
        password = request.form['password']
        edsl_api_key = request.form.get('edsl_api_key', '')
        
        # Check if user already exists
        if User.query.filter_by(username=username).first():
            flash('Username already exists', 'error')
            return render_template('login.html')
        
        if User.query.filter_by(email=email).first():
            flash('Email already registered', 'error')
            return render_template('login.html')
        
        # Create new user
        user = User(
            username=username,
            email=email,
            password_hash=generate_password_hash(password),
            edsl_api_key=edsl_api_key
        )
        
        db.session.add(user)
        db.session.commit()
        
        session['user_id'] = user.id
        flash('Registration successful!', 'success')
        return redirect(url_for('index'))
    
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.pop('user_id', None)
    flash('You have been logged out', 'info')
    return redirect(url_for('index'))

@app.route('/upload', methods=['GET', 'POST'])
def upload():
    if 'user_id' not in session:
        flash('Please log in to upload papers', 'error')
        return redirect(url_for('login'))
    
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('No file selected', 'error')
            return redirect(request.url)
        
        file = request.files['file']
        title = request.form.get('title', '').strip()
        
        if file.filename == '':
            flash('No file selected', 'error')
            return redirect(request.url)
        
        if not title:
            flash('Please provide a title for your paper', 'error')
            return redirect(request.url)
        
        if file and allowed_file(file.filename):
            # Generate unique filename
            filename = secure_filename(file.filename)
            unique_filename = f"{uuid.uuid4()}_{filename}"
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(filepath)
            
            # Create review session
            review_session = ReviewSession(
                title=title,
                paper_filename=filename,
                paper_path=filepath,
                user_id=session['user_id']
            )
            
            db.session.add(review_session)
            db.session.commit()
            
            flash('Paper uploaded successfully! Starting review process...', 'success')
            return redirect(url_for('review_process', session_id=review_session.id))
        else:
            flash('Invalid file type. Please upload a PDF file.', 'error')
    
    return render_template('upload.html')

@app.route('/review/<int:session_id>')
def review_process(session_id):
    if 'user_id' not in session:
        flash('Please log in to view review sessions', 'error')
        return redirect(url_for('login'))
    
    review_session = ReviewSession.query.filter_by(id=session_id, user_id=session['user_id']).first()
    if not review_session:
        flash('Review session not found', 'error')
        return redirect(url_for('index'))
    
    return render_template('review_process.html', session=review_session)

@app.route('/start_review/<int:session_id>')
def start_review(session_id):
    if 'user_id' not in session:
        return jsonify({'error': 'Not authenticated'}), 401
    
    review_session = ReviewSession.query.filter_by(id=session_id, user_id=session['user_id']).first()
    if not review_session:
        return jsonify({'error': 'Review session not found'}), 404
    
    user = User.query.get(session['user_id'])
    
    try:
        # Initialize EDSL service
        edsl_service = EDSLService(user.edsl_api_key)
        
        # Update session status
        review_session.status = 'processing'
        review_session.processing_stage = 'referee_reports'
        db.session.commit()
        
        # Generate referee reports
        logger.info(f"Starting referee report generation for session {session_id}")
        referee_reports = edsl_service.generate_referee_reports(review_session.paper_path)
        review_session.referee_reports = referee_reports
        review_session.processing_stage = 'author_response'
        db.session.commit()
        
        # Generate author response
        logger.info(f"Starting author response generation for session {session_id}")
        author_response = edsl_service.generate_author_response(review_session.paper_path, referee_reports)
        review_session.author_response = author_response
        review_session.processing_stage = 'referee_rebuttals'
        db.session.commit()
        
        # Generate referee rebuttals
        logger.info(f"Starting referee rebuttal generation for session {session_id}")
        referee_rebuttals = edsl_service.generate_referee_rebuttals(referee_reports, author_response)
        review_session.referee_rebuttals = referee_rebuttals
        
        # Mark as completed
        review_session.status = 'completed'
        review_session.completed_at = datetime.utcnow()
        review_session.processing_stage = 'completed'
        db.session.commit()
        
        logger.info(f"Review process completed for session {session_id}")
        return jsonify({'status': 'completed', 'message': 'Review process completed successfully!'})
        
    except Exception as e:
        logger.error(f"Error during review process for session {session_id}: {str(e)}")
        review_session.status = 'error'
        review_session.error_message = str(e)
        db.session.commit()
        return jsonify({'error': f'Review process failed: {str(e)}'}), 500

@app.route('/review_status/<int:session_id>')
def review_status(session_id):
    if 'user_id' not in session:
        return jsonify({'error': 'Not authenticated'}), 401
    
    review_session = ReviewSession.query.filter_by(id=session_id, user_id=session['user_id']).first()
    if not review_session:
        return jsonify({'error': 'Review session not found'}), 404
    
    return jsonify({
        'status': review_session.status,
        'processing_stage': review_session.processing_stage,
        'error_message': review_session.error_message
    })

@app.route('/results/<int:session_id>')
def results(session_id):
    if 'user_id' not in session:
        flash('Please log in to view results', 'error')
        return redirect(url_for('login'))
    
    review_session = ReviewSession.query.filter_by(id=session_id, user_id=session['user_id']).first()
    if not review_session:
        flash('Review session not found', 'error')
        return redirect(url_for('index'))
    
    if review_session.status != 'completed':
        flash('Review process is not completed yet', 'warning')
        return redirect(url_for('review_process', session_id=session_id))
    
    return render_template('results.html', session=review_session)

@app.route('/download/<int:session_id>')
def download_results(session_id):
    if 'user_id' not in session:
        flash('Please log in to download results', 'error')
        return redirect(url_for('login'))
    
    review_session = ReviewSession.query.filter_by(id=session_id, user_id=session['user_id']).first()
    if not review_session or review_session.status != 'completed':
        flash('Results not available for download', 'error')
        return redirect(url_for('index'))
    
    # Create a temporary zip file with all results
    with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as tmp_file:
        with zipfile.ZipFile(tmp_file.name, 'w') as zip_file:
            # Add referee reports
            if review_session.referee_reports:
                zip_file.writestr('referee_reports.txt', review_session.referee_reports)
            
            # Add author response
            if review_session.author_response:
                zip_file.writestr('author_response.txt', review_session.author_response)
            
            # Add referee rebuttals
            if review_session.referee_rebuttals:
                zip_file.writestr('referee_rebuttals.txt', review_session.referee_rebuttals)
        
        return send_file(
            tmp_file.name,
            as_attachment=True,
            download_name=f"{review_session.title}_review_results.zip",
            mimetype='application/zip'
        )

@app.route('/sessions')
def sessions():
    if 'user_id' not in session:
        flash('Please log in to view your sessions', 'error')
        return redirect(url_for('login'))
    
    user_sessions = ReviewSession.query.filter_by(user_id=session['user_id']).order_by(ReviewSession.created_at.desc()).all()
    return render_template('sessions.html', sessions=user_sessions)