JARVIS/backend/app/routers/folder.py

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession
from typing import List
import shutil

from app.database import get_db
from app.models.folder import Folder
from app.models.user import User
from app.routers.auth import get_current_user
from app.schemas.folder import FolderCreate, FolderOut, FolderTreeOut, FolderUpdate
from app.services.document_service import DocumentService

router = APIRouter(prefix="/api/folders", tags=["文件夹"])


def build_folder_tree(folders: list[Folder], parent_id: str = None) -> List[FolderTreeOut]:
    tree = []
    for folder in folders:
        if folder.parent_id == parent_id:
            children = build_folder_tree(folders, folder.id)
            tree.append(FolderTreeOut(
                id=folder.id,
                name=folder.name,
                parent_id=folder.parent_id,
                children=children,
            ))
    return tree


@router.get("", response_model=List[FolderTreeOut])
async def get_folders(
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    result = await db.execute(
        select(Folder).where(Folder.user_id == current_user.id)
    )
    folders = result.scalars().all()
    return build_folder_tree(list(folders))


@router.post("", response_model=FolderOut, status_code=status.HTTP_201_CREATED)
async def create_folder(
    folder_data: FolderCreate,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    if folder_data.parent_id:
        result = await db.execute(
            select(Folder).where(
                and_(Folder.id == folder_data.parent_id, Folder.user_id == current_user.id)
            )
        )
        if not result.scalar_one_or_none():
            raise HTTPException(status_code=404, detail="父文件夹不存在")

    result = await db.execute(
        select(Folder).where(
            and_(
                Folder.user_id == current_user.id,
                Folder.parent_id == folder_data.parent_id,
                Folder.name == folder_data.name,
            )
        )
    )
    if result.scalar_one_or_none():
        raise HTTPException(status_code=400, detail="同名文件夹已存在")

    folder = Folder(
        user_id=current_user.id,
        name=folder_data.name,
        parent_id=folder_data.parent_id,
    )
    db.add(folder)
    await db.commit()
    await db.refresh(folder)

    document_service = DocumentService(db, current_user.id)
    await document_service.ensure_folder_directory(current_user.id, folder.id)
    return folder


@router.put("/{folder_id}", response_model=FolderOut)
async def rename_folder(
    folder_id: str,
    folder_data: FolderUpdate,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    result = await db.execute(
        select(Folder).where(
            and_(Folder.id == folder_id, Folder.user_id == current_user.id)
        )
    )
    folder = result.scalar_one_or_none()
    if not folder:
        raise HTTPException(status_code=404, detail="文件夹不存在")

    old_name = folder.name
    folder.name = folder_data.name

    document_service = DocumentService(db, current_user.id)
    await document_service.rename_folder_directory(current_user.id, folder.id, old_name, folder_data.name)
    await db.commit()
    await db.refresh(folder)
    return folder


@router.delete("/{folder_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_folder(
    folder_id: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    from app.models.document import Document
    from app.services.knowledge_service import KnowledgeService

    result = await db.execute(
        select(Folder).where(
            and_(Folder.id == folder_id, Folder.user_id == current_user.id)
        )
    )
    folder = result.scalar_one_or_none()
    if not folder:
        raise HTTPException(status_code=404, detail="文件夹不存在")

    document_service = DocumentService(db, current_user.id)
    folder_path = await document_service._get_storage_directory(current_user.id, folder_id)

    async def delete_recursive(fid: str):
        children = await db.execute(
            select(Folder).where(Folder.parent_id == fid)
        )
        for child in children.scalars():
            await delete_recursive(child.id)

        docs = await db.execute(
            select(Document).where(Document.folder_id == fid)
        )
        for doc in docs.scalars():
            knowledge_service = KnowledgeService(db, current_user.id)
            await knowledge_service.delete_from_vectorstore(current_user.id, doc.id)
            await db.delete(doc)

        folder_to_delete = await db.get(Folder, fid)
        if folder_to_delete:
            await db.delete(folder_to_delete)

    await delete_recursive(folder_id)
    await db.commit()

    if folder_path.exists():
        shutil.rmtree(folder_path, ignore_errors=True)