@app.route('/document/<int:document_id>/edit', methods=['GET', 'POST'])
@login_required
def edit_document(document_id):
    conn = get_db_connection()
    
    # Get document
    document = conn.execute('SELECT * FROM documents WHERE id = ?', (document_id,)).fetchone()
    
    if not document:
        conn.close()
        flash('Document not found', 'error')
        return redirect(url_for('files'))
    
    # Check ownership
    if document['user_id'] != session['user_id'] and session['role'] != 'admin':
        conn.close()
        flash('You do not have permission to edit this document', 'error')
        return redirect(url_for('files'))
    
    # Get available categories
    categories = ['admin', 'accounting', 'hr', 'marketing', 'legal', 'general', 'other']
    
    if request.method == 'POST':
        custom_filename = request.form.get('custom_filename', '').strip()
        category = request.form.get('category', 'general')
        visibility = request.form.get('visibility', 'private')
        
        # Use custom filename if provided, otherwise keep existing
        if not custom_filename:
            custom_filename = document['custom_filename']
        
        # Update document
        conn.execute('''
            UPDATE documents 
            SET custom_filename = ?, category = ?, visibility = ?, updated_at = CURRENT_TIMESTAMP
            WHERE id = ?
        ''', (custom_filename, category, visibility, document_id))
        
        conn.commit()
        conn.close()
        
        flash('Document updated successfully', 'success')
        return redirect(url_for('view_document', document_id=document_id))
    
    conn.close()
    
    return render_template('edit_document.html', document=document, categories=categories)