flask-vim-docs/app/routes.py

from flask import Blueprint, render_template, request, redirect, url_for, flash, jsonify, send_file
from flask_login import current_user, login_required
from app import db
from app.models.document import Document, Category, Tag
from app.models.user import User
import json
from datetime import datetime
import io
from sqlalchemy import or_

main = Blueprint('main', __name__)

# Helper function to convert categories to a hierarchical dictionary
def build_category_hierarchy(categories):
    categories_dict = []
    
    # Start with root categories (parent_id is None)
    root_categories = [c for c in categories if c.parent_id is None]
    
    # Helper function to convert category and its children to dict
    def category_to_dict(category):
        cat_dict = {
            'id': category.id,
            'name': category.name,
            'icon': category.icon,
            'is_root': getattr(category, 'is_root', False),
            'parent_id': category.parent_id,
            'children': []
        }
        
        # Find and add children
        for child in [c for c in categories if c.parent_id == category.id]:
            cat_dict['children'].append(category_to_dict(child))
            
        return cat_dict
    
    # Convert all root categories and their children
    for category in root_categories:
        categories_dict.append(category_to_dict(category))
    
    return categories_dict

@main.route('/')
@login_required
def index():
    """Home page showing categories and recent documents"""
    root_categories = Category.query.filter_by(
        user_id=current_user.id, 
        parent_id=None
    ).all()
    
    recent_docs = Document.query.filter_by(
        user_id=current_user.id
    ).order_by(Document.updated_date.desc()).limit(5).all()
    
    return render_template('index.html', categories=root_categories, recent_docs=recent_docs)

@main.route('/category/<int:category_id>')
@login_required
def view_category(category_id):
    """View a specific category and its documents"""
    category = Category.query.filter_by(id=category_id, user_id=current_user.id).first_or_404()
    return render_template('category.html', category=category)

@main.route('/category/<int:category_id>/all_documents')
@login_required
def category_all_documents(category_id):
    """View all documents in a category, including those in subcategories"""
    category = Category.query.filter_by(id=category_id, user_id=current_user.id).first_or_404()
    
    # Get all documents from this category
    documents = []
    
    # Add documents from the current category
    for doc in category.documents:
        documents.append(doc)
    
    # Helper function to recursively collect documents from subcategories
    def collect_documents(cat):
        docs = []
        for doc in cat.documents:
            docs.append(doc)
        
        for child in cat.children:
            docs.extend(collect_documents(child))
        
        return docs
    
    # Collect documents from all subcategories
    for subcategory in category.children:
        documents.extend(collect_documents(subcategory))
    
    # Sort by updated_date, with most recent first
    documents.sort(key=lambda x: x.updated_date, reverse=True)
    
    return render_template('category_documents.html', category=category, documents=documents)

@main.route('/document/<int:doc_id>', methods=['GET'])
@login_required
def view_document(doc_id):
    """View a document in read mode"""
    document = Document.query.filter_by(id=doc_id, user_id=current_user.id).first_or_404()
    return render_template('document_view.html', document=document)

@main.route('/document/<int:doc_id>/edit', methods=['GET'])
@login_required
def edit_document(doc_id):
    """Edit a document with the Vim editor"""
    document = Document.query.filter_by(id=doc_id, user_id=current_user.id).first_or_404()
    categories = Category.query.filter_by(user_id=current_user.id).all()
    tags = Tag.query.filter_by(user_id=current_user.id).all()
    
    # Convert categories to dictionaries for JSON serialization
    categories_dict = build_category_hierarchy(categories)
    
    return render_template('document_edit.html', document=document, categories=categories_dict, tags=tags)

@main.route('/document/new', methods=['GET'])
@login_required
def new_document():
    """Create a new document with the Vim editor"""
    categories = Category.query.filter_by(user_id=current_user.id).all()
    tags = Tag.query.filter_by(user_id=current_user.id).all()
    category_id = request.args.get('category')
    
    document = None
    preselected_category_id = None
    
    # Find the root category
    root_category = Category.query.filter_by(user_id=current_user.id, is_root=True).first()
    
    if category_id:
        # If a specific category was requested, use that
        category = Category.query.filter_by(id=category_id, user_id=current_user.id).first()
        if category:
            preselected_category_id = int(category_id)
    elif root_category:
        # Otherwise default to the root category
        preselected_category_id = root_category.id
    
    # Create a blank document
    document = Document(
        title="Untitled Document", 
        content="", 
        category_id=preselected_category_id,
        user_id=current_user.id
    )
    
    # Convert categories to dictionaries for JSON serialization
    categories_dict = build_category_hierarchy(categories)
    
    return render_template('document_edit.html', document=document, categories=categories_dict, tags=tags, preselected_category_id=preselected_category_id)

@main.route('/api/document', methods=['POST'])
@login_required
def save_document():
    """Save a document (new or existing)"""
    data = request.json
    
    # Find the root category to use as default
    root_category = Category.query.filter_by(user_id=current_user.id, is_root=True).first()
    
    # Ensure we have a root category
    if not root_category:
        # Create a root category if it doesn't exist
        root_category = Category(
            name='Library',
            icon='mdi-bookshelf',
            description='General storage for documents and categories',
            user_id=current_user.id,
            is_root=True
        )
        db.session.add(root_category)
        db.session.flush()  # Get the ID without committing yet
    
    # All documents must have a category - defaults to root
    if 'id' in data and data['id']:
        # Update existing document - verify ownership
        document = Document.query.filter_by(id=data['id'], user_id=current_user.id).first_or_404()
        document.title = data['title']
        document.content = data['content']
        document.category_id = data['category_id'] if data['category_id'] else root_category.id
    else:
        # Create new document
        document = Document(
            title=data['title'],
            content=data['content'],
            category_id=data['category_id'] if data['category_id'] else root_category.id,
            user_id=current_user.id
        )
        db.session.add(document)
    
    # Handle tags
    if 'tags' in data:
        # Clear existing tags
        document.tags = []
        
        # Process tags, preventing duplicates
        processed_tag_names = set()
        for tag_name in data['tags']:
            # Skip if we've already added this tag (case insensitive)
            if tag_name.lower() in processed_tag_names:
                continue
                
            # Add to our processed set to prevent duplicates in this request
            processed_tag_names.add(tag_name.lower())
            
            # Check if tag already exists for this user
            tag = Tag.query.filter(
                Tag.name.ilike(tag_name),
                Tag.user_id == current_user.id
            ).first()
            
            if not tag:
                # Create new tag
                tag = Tag(name=tag_name, user_id=current_user.id)
                db.session.add(tag)
                
            # Add tag to document
            document.tags.append(tag)
    
    db.session.commit()
    return jsonify(document.to_dict())

@main.route('/api/document/<int:doc_id>', methods=['DELETE'])
@login_required
def delete_document(doc_id):
    """Delete a document"""
    document = Document.query.filter_by(id=doc_id, user_id=current_user.id).first_or_404()
    db.session.delete(document)
    db.session.commit()
    return jsonify({'success': True})

@main.route('/api/document/<int:doc_id>', methods=['GET'])
@login_required
def get_document(doc_id):
    """Get a document by ID"""
    document = Document.query.filter_by(id=doc_id, user_id=current_user.id).first_or_404()
    return jsonify(document.to_dict())

@main.route('/api/document/<int:doc_id>', methods=['PATCH'])
@login_required
def update_document(doc_id):
    """Update a document's properties (like category)"""
    document = Document.query.filter_by(id=doc_id, user_id=current_user.id).first_or_404()
    data = request.json
    
    # Update category if provided
    if 'category_id' in data and data['category_id']:
        # Verify the category exists and belongs to user
        category = Category.query.filter_by(id=data['category_id'], user_id=current_user.id).first_or_404()
        document.category_id = category.id
    
    db.session.commit()
    return jsonify(document.to_dict())

@main.route('/api/category', methods=['POST'])
@login_required
def save_category():
    """Save a category (new or existing)"""
    data = request.json
    
    if 'id' in data and data['id']:
        # Update existing category - verify ownership
        category = Category.query.filter_by(id=data['id'], user_id=current_user.id).first_or_404()
        category.name = data['name']
        category.icon = data['icon']
        category.description = data.get('description', '')
        
        # Only change parent if it belongs to the same user
        if data.get('parent_id'):
            parent = Category.query.filter_by(id=data['parent_id'], user_id=current_user.id).first()
            if parent:
                category.parent_id = parent.id
    else:
        # Create new category
        category = Category(
            name=data['name'],
            icon=data['icon'],
            description=data.get('description', ''),
            parent_id=data['parent_id'] if data.get('parent_id') else None,
            user_id=current_user.id
        )
        db.session.add(category)
    
    db.session.commit()
    return jsonify(category.to_dict())

@main.route('/api/category/<int:category_id>', methods=['DELETE'])
@login_required
def delete_category(category_id):
    """Delete a category and optionally reassign documents"""
    category = Category.query.filter_by(id=category_id, user_id=current_user.id).first_or_404()
    
    # Can't delete root category
    if category.is_root:
        return jsonify({'error': 'Cannot delete root category'}), 400
    
    # Get target category for documents if specified
    new_category_id = request.args.get('new_category_id')
    if new_category_id:
        new_category = Category.query.filter_by(id=new_category_id, user_id=current_user.id).first()
        if new_category:
            # Reassign documents
            for doc in category.documents:
                doc.category_id = new_category.id
    else:
        # Move documents to no category
        for doc in category.documents:
            doc.category_id = None
    
    # Also handle child categories
    for child in category.children:
        if new_category_id:
            child.parent_id = new_category_id
        else:
            child.parent_id = None
    
    db.session.delete(category)
    db.session.commit()
    return jsonify({'success': True})

@main.route('/api/search', methods=['GET'])
@login_required
def search_documents():
    """Search for documents by title, content, or tags"""
    query = request.args.get('q', '')
    
    if not query or len(query) < 2:
        return jsonify({'results': []})
    
    # Search in title, content, and tags for current user's documents only
    docs = Document.query.filter(
        Document.user_id == current_user.id,
        or_(
            Document.title.ilike(f'%{query}%'),
            Document.content.ilike(f'%{query}%'),
            Document.tags.any(Tag.name.ilike(f'%{query}%'))
        )
    ).limit(10).all()
    
    results = []
    for doc in docs:
        # Find match in content
        match = None
        if query.lower() in doc.content.lower():
            # Find the sentence containing the match
            content_lower = doc.content.lower()
            query_pos = content_lower.find(query.lower())
            
            # Get a snippet around the match
            start = max(0, content_lower.rfind('.', 0, query_pos) + 1)
            end = content_lower.find('.', query_pos)
            if end == -1:
                end = min(len(doc.content), query_pos + 200)
            
            match = doc.content[start:end].strip()
        
        # Get category name
        category_name = doc.category.name if doc.category else None
        
        results.append({
            'id': doc.id,
            'title': doc.title,
            'category': category_name,
            'tags': [tag.name for tag in doc.tags],
            'match': match
        })
    
    return jsonify({'results': results})

@main.route('/api/tags', methods=['GET'])
@login_required
def get_tags():
    """Get all tags for the current user"""
    tags = Tag.query.filter_by(user_id=current_user.id).all()
    return jsonify({'tags': [tag.to_dict() for tag in tags]})

@main.route('/document/<int:doc_id>/export', methods=['GET'])
@login_required
def export_document(doc_id):
    """Export a document as a markdown file"""
    document = Document.query.filter_by(id=doc_id, user_id=current_user.id).first_or_404()
    
    # Create a file-like object in memory
    file_data = io.BytesIO(document.content.encode('utf-8'))
    file_data.seek(0)
    
    return send_file(
        file_data,
        mimetype='text/markdown',
        as_attachment=True,
        download_name=f"{document.title.replace(' ', '_')}.md"
    )

@main.route('/api/categories', methods=['GET'])
@login_required
def get_categories():
    """Get all root categories with their children for the current user"""
    root_categories = Category.query.filter_by(
        user_id=current_user.id, 
        parent_id=None
    ).all()
    return jsonify([category.to_dict() for category in root_categories])

@main.route('/category/new', methods=['GET'])
@login_required
def new_category():
    """Create a new category or subcategory"""
    parent_id = request.args.get('parent_id')
    parent = None
    
    if parent_id:
        parent = Category.query.filter_by(id=parent_id, user_id=current_user.id).first_or_404()
    
    return render_template('category_edit.html', category=None, parent=parent)

@main.route('/category/<int:category_id>/edit', methods=['GET'])
@login_required
def edit_category(category_id):
    """Edit an existing category"""
    category = Category.query.filter_by(id=category_id, user_id=current_user.id).first_or_404()
    return render_template('category_edit.html', category=category, parent=category.parent)