video-query/backend/chunked_upload.py

import os
import uuid
import json
from flask import Blueprint, request, jsonify, current_app
from werkzeug.utils import secure_filename
import logging
from auth import lenient_auth
from system_utils import system_utils
from error_reporter import ErrorReporter, ErrorCategory

logger = logging.getLogger('video_query')

# Create blueprint for handling chunked uploads
chunked_upload_bp = Blueprint('chunked_upload', __name__)

# CORS preflight handler for all blueprint routes
def handle_cors_preflight():
    """Handle CORS preflight requests"""
    origin = request.headers.get('Origin')
    allowed_origins = ['https://brandtechsandbox.oliver.solutions', 'http://localhost:3000']

    response = jsonify({})

    # Allow the origin if it's in our allowed list
    if origin in allowed_origins:
        response.headers.add('Access-Control-Allow-Origin', origin)
    else:
        # Default to production origin
        response.headers.add('Access-Control-Allow-Origin', 'https://brandtechsandbox.oliver.solutions')

    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization,X-Requested-With')
    response.headers.add('Access-Control-Allow-Methods', 'GET,POST,OPTIONS')
    response.headers.add('Access-Control-Max-Age', '86400')  # 24 hours
    response.headers.add('Access-Control-Allow-Credentials', 'true')
    return response

# Track upload sessions
active_uploads = {}

@chunked_upload_bp.route('/api/init-upload', methods=['OPTIONS'])
def init_upload_options():
    """Handle CORS preflight for init-upload"""
    return handle_cors_preflight()

@chunked_upload_bp.route('/api/init-upload', methods=['POST'])
@lenient_auth
def init_upload():
    """Initialize a new chunked upload session"""
    if not request.is_json:
        return jsonify({"success": False, "message": "Request must be JSON"}), 400
    
    data = request.get_json()
    filename = data.get('filename')
    total_size = data.get('size')
    
    if not filename or not total_size:
        return jsonify({"success": False, "message": "Filename and size are required"}), 400
    
    # Generate a unique ID for this upload
    upload_id = str(uuid.uuid4())
    
    # Generate a unique filename
    original_filename = secure_filename(filename)
    unique_filename = f"{upload_id}_{original_filename}"
    upload_path = os.path.join(current_app.config['UPLOAD_FOLDER'], unique_filename)
    
    # Create/ensure the upload folder exists
    os.makedirs(current_app.config['UPLOAD_FOLDER'], exist_ok=True)
    
    # Initialize an empty file
    with open(upload_path, 'wb') as f:
        pass
    
    # Store upload info
    active_uploads[upload_id] = {
        'path': upload_path,
        'original_filename': original_filename,
        'total_size': total_size,
        'uploaded_size': 0,
        'complete': False
    }
    
    logger.info(f"Initialized upload session {upload_id} for {filename} ({total_size} bytes)")
    
    return jsonify({
        "success": True,
        "upload_id": upload_id
    })

@chunked_upload_bp.route('/api/upload-chunk/<upload_id>', methods=['OPTIONS'])
def upload_chunk_options(upload_id):
    """Handle CORS preflight for upload-chunk"""
    return handle_cors_preflight()

@chunked_upload_bp.route('/api/upload-chunk/<upload_id>', methods=['POST'])
@lenient_auth
def upload_chunk(upload_id):
    """Handle a chunk of file data"""
    if upload_id not in active_uploads:
        return jsonify({"success": False, "message": "Invalid upload ID"}), 400
    
    upload = active_uploads[upload_id]
    
    # Check if we have the file chunk
    if 'chunk' not in request.files:
        return jsonify({"success": False, "message": "No chunk in request"}), 400
    
    chunk = request.files['chunk']
    chunk_number = request.form.get('chunk_number', 0, type=int)
    
    # Update the file with this chunk
    with open(upload['path'], 'ab') as f:
        chunk_data = chunk.read()
        chunk_size = len(chunk_data)
        f.write(chunk_data)
    
    # Update upload state
    upload['uploaded_size'] += chunk_size
    progress = min(100, round((upload['uploaded_size'] / upload['total_size']) * 100))
    
    logger.info(f"Received chunk {chunk_number} for upload {upload_id} - Progress: {progress}%")
    
    # Check if upload is complete
    if upload['uploaded_size'] >= upload['total_size']:
        upload['complete'] = True
        logger.info(f"Upload complete for {upload_id}")
    
    return jsonify({
        "success": True,
        "upload_id": upload_id,
        "chunk_number": chunk_number,
        "bytes_received": chunk_size,
        "total_received": upload['uploaded_size'],
        "progress": progress,
        "complete": upload['complete']
    })

@chunked_upload_bp.route('/api/complete-upload/<upload_id>', methods=['OPTIONS'])
def complete_upload_options(upload_id):
    """Handle CORS preflight for complete-upload"""
    return handle_cors_preflight()

@chunked_upload_bp.route('/api/complete-upload/<upload_id>', methods=['POST'])
@lenient_auth
def complete_upload(upload_id):
    """Mark an upload as complete and return the file path for processing"""
    if upload_id not in active_uploads:
        return jsonify({"success": False, "message": "Invalid upload ID"}), 400

    upload = active_uploads[upload_id]

    # Verify the upload is actually complete
    if not upload['complete']:
        # Check the file size
        if os.path.exists(upload['path']):
            actual_size = os.path.getsize(upload['path'])
            if actual_size >= upload['total_size']:
                upload['complete'] = True
                upload['uploaded_size'] = actual_size
                logger.info(f"Manually verified upload complete for {upload_id}")
            else:
                logger.warning(f"Upload not complete for {upload_id}: {actual_size}/{upload['total_size']} bytes")
                return jsonify({
                    "success": False,
                    "message": f"Upload not complete: {actual_size}/{upload['total_size']} bytes"
                }), 400
        else:
            logger.error(f"Upload file not found for {upload_id}")
            return jsonify({"success": False, "message": "Upload file not found"}), 500

    # Additional validation: Verify file integrity with FFprobe
    import subprocess
    import time

    # Wait a moment for file system to fully sync the file
    time.sleep(0.5)

    try:
        logger.info(f"Validating uploaded file integrity for {upload_id}")
        ffprobe_path = system_utils.find_ffprobe()
        probe_result = subprocess.run(
            [ffprobe_path, '-v', 'error', '-show_entries', 'format=duration,format_name',
             '-of', 'default=noprint_wrappers=1', upload['path']],
            capture_output=True, text=True, timeout=15
        )

        if probe_result.returncode != 0:
            error_detail = probe_result.stderr.strip()
            logger.error(f"Upload validation failed for {upload_id}: {error_detail}")

            # Check for specific "moov atom not found" error
            if "moov atom not found" in error_detail.lower():
                error_msg = "Video file is incomplete or corrupted (missing header). The upload may have been interrupted. Please try uploading again."
            else:
                error_msg = f"Video file validation failed: {error_detail}"

            # Delete the invalid file
            try:
                os.remove(upload['path'])
                logger.info(f"Deleted invalid upload file for {upload_id}")
            except Exception as del_err:
                logger.warning(f"Could not delete invalid file: {str(del_err)}")

            # Remove from active uploads
            del active_uploads[upload_id]

            return jsonify({
                "success": False,
                "message": error_msg
            }), 400

        logger.info(f"Upload validation successful for {upload_id}")

    except subprocess.TimeoutExpired:
        logger.warning(f"Upload validation timed out for {upload_id} - proceeding anyway")
    except FileNotFoundError as e:
        error_report = ErrorReporter.capture_error(
            e,
            category=ErrorCategory.SYSTEM_ERROR,
            context={'upload_id': upload_id, 'operation': 'upload_validation'},
            severity='warning'
        )
        logger.warning(f"ffprobe not found - skipping upload validation for {upload_id}")
    except Exception as val_err:
        error_report = ErrorReporter.capture_error(
            val_err,
            category=ErrorCategory.UPLOAD_ERROR,
            context={'upload_id': upload_id, 'operation': 'upload_validation'},
            severity='warning'
        )
        logger.warning(f"Error during upload validation for {upload_id}: {str(val_err)} - proceeding anyway")

    logger.info(f"Upload {upload_id} marked as complete: {upload['original_filename']}")

    return jsonify({
        "success": True,
        "upload_id": upload_id,
        "file_path": upload['path'],
        "filename": upload['original_filename'],
        "size": upload['uploaded_size']
    })

@chunked_upload_bp.route('/api/cancel-upload/<upload_id>', methods=['OPTIONS'])
def cancel_upload_options(upload_id):
    """Handle CORS preflight for cancel-upload"""
    return handle_cors_preflight()

@chunked_upload_bp.route('/api/cancel-upload/<upload_id>', methods=['POST'])
@lenient_auth
def cancel_upload(upload_id):
    """Cancel an upload and delete the partial file"""
    if upload_id not in active_uploads:
        return jsonify({"success": False, "message": "Invalid upload ID"}), 400
    
    upload = active_uploads[upload_id]
    
    # Delete the partial file
    if os.path.exists(upload['path']):
        try:
            os.remove(upload['path'])
            logger.info(f"Deleted partial upload for {upload_id}")
        except Exception as e:
            logger.error(f"Error deleting partial upload for {upload_id}: {str(e)}")
    
    # Remove from active uploads
    del active_uploads[upload_id]
    
    return jsonify({
        "success": True,
        "message": "Upload cancelled"
    })