AgentEvalTool/backend/agenteval/web/routers/files.py

"""API routes for file management (categories + file upload/download)."""

import mimetypes
import uuid
from pathlib import Path

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from pydantic import BaseModel
from sqlmodel import Session

from agenteval.config import get_settings
from agenteval.storage.db import FILES_DIR
from agenteval.storage.file_repository import FileCategoryRepository, FileRecordRepository
from agenteval.web.deps import get_db

router = APIRouter()

# ── helpers ────────────────────────────────────────────────────────────


def _get_allowed_extensions() -> set[str]:
    """Return the set of allowed lowercase extensions from settings."""
    raw = get_settings().allowed_extensions
    return {ext.strip().lower() for ext in raw.split(",") if ext.strip()}


def _get_max_bytes() -> int:
    return get_settings().max_upload_size_mb * 1024 * 1024


def _validate_extension(filename: str) -> str:
    """Validate the file extension and return the lowercase extension."""
    ext = Path(filename).suffix.lstrip(".").lower()
    if not ext:
        raise HTTPException(status_code=400, detail="文件没有扩展名")
    allowed = _get_allowed_extensions()
    if ext not in allowed:
        raise HTTPException(
            status_code=400,
            detail=f"不支持的文件类型 .{ext}，允许的类型: {', '.join(sorted(allowed))}",
        )
    return ext


# ── API models ─────────────────────────────────────────────────────────


class CategoryCreate(BaseModel):
    name: str
    parent_id: str | None = None


class CategoryUpdate(BaseModel):
    name: str


# ── Category endpoints ─────────────────────────────────────────────────


@router.get("/categories")
def list_categories(session: Session = Depends(get_db)) -> list[dict]:
    """List categories as a nested tree."""
    return FileCategoryRepository(session).get_tree()


@router.post("/categories")
def create_category(body: CategoryCreate, session: Session = Depends(get_db)) -> dict:
    if not body.name.strip():
        raise HTTPException(status_code=400, detail="分类名称不能为空")
    cat = FileCategoryRepository(session).create(
        name=body.name.strip(),
        parent_id=body.parent_id,
    )
    return {
        "key": cat.id,
        "title": cat.name,
        "parent_id": cat.parent_id,
        "children": [],
    }


@router.put("/categories/{category_id}")
def update_category(category_id: str, body: CategoryUpdate, session: Session = Depends(get_db)) -> dict:
    if not body.name.strip():
        raise HTTPException(status_code=400, detail="分类名称不能为空")
    cat = FileCategoryRepository(session).update(category_id, body.name.strip())
    if not cat:
        raise HTTPException(status_code=404, detail="分类不存在")
    return {"ok": True, "name": cat.name}


@router.delete("/categories/{category_id}")
def delete_category(category_id: str, session: Session = Depends(get_db)) -> dict:
    if not FileCategoryRepository(session).delete(category_id):
        raise HTTPException(status_code=404, detail="分类不存在")
    return {"ok": True}


# ── File endpoints ─────────────────────────────────────────────────────


@router.get("")
def list_files(category_id: str | None = None, session: Session = Depends(get_db)) -> list[dict]:
    """List file records, optionally filtered by category."""
    records = FileRecordRepository(session).list_all(category_id=category_id)
    return [
        {
            "id": r.id,
            "original_name": r.original_name,
            "file_size": r.file_size,
            "mime_type": r.mime_type,
            "file_ext": r.file_ext,
            "category_id": r.category_id,
            "created_at": r.created_at.isoformat() if r.created_at else None,
        }
        for r in records
    ]


@router.post("/upload")
async def upload_file(
    file: UploadFile = File(...),
    category_id: str | None = Form(default=None),
    session: Session = Depends(get_db),
) -> dict:
    """Upload a single file."""
    if not file.filename:
        raise HTTPException(status_code=400, detail="文件名为空")

    ext = _validate_extension(file.filename)

    # Read content and validate size
    content = await file.read()
    max_bytes = _get_max_bytes()
    if len(content) > max_bytes:
        raise HTTPException(
            status_code=413,
            detail=f"文件大小超过限制 ({get_settings().max_upload_size_mb}MB)",
        )

    # Generate storage name
    storage_name = f"{uuid.uuid4()}.{ext}"

    # Guess MIME type
    mime_type, _ = mimetypes.guess_type(file.filename)
    if not mime_type:
        mime_type = "application/octet-stream"

    # Validate category exists if provided
    if category_id:
        cat = FileCategoryRepository(session).get(category_id)
        if not cat:
            raise HTTPException(status_code=400, detail="分类不存在")

    # Write physical file
    file_path = FILES_DIR / storage_name
    file_path.write_bytes(content)

    # Create DB record
    record = FileRecordRepository(session).create(
        original_name=file.filename,
        storage_name=storage_name,
        file_size=len(content),
        mime_type=mime_type,
        file_ext=ext,
        category_id=category_id if category_id else None,
    )

    return {
        "id": record.id,
        "original_name": record.original_name,
        "file_size": record.file_size,
        "mime_type": record.mime_type,
        "file_ext": record.file_ext,
        "category_id": record.category_id,
        "created_at": record.created_at.isoformat() if record.created_at else None,
    }


@router.get("/{file_id}/download")
def download_file(file_id: str, session: Session = Depends(get_db)):
    """Download a file by its record ID."""
    record = FileRecordRepository(session).get(file_id)
    if not record:
        raise HTTPException(status_code=404, detail="文件不存在")

    file_path = FILES_DIR / record.storage_name
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="物理文件不存在")

    return FileResponse(
        path=str(file_path),
        filename=record.original_name,
        media_type=record.mime_type or "application/octet-stream",
    )


@router.delete("/{file_id}")
def delete_file(file_id: str, session: Session = Depends(get_db)) -> dict:
    if not FileRecordRepository(session).delete(file_id):
        raise HTTPException(status_code=404, detail="文件不存在")
    return {"ok": True}