audio2text/app/controllers/task_router.py

"""任务路由：列表 / 状态 / 下载字幕 / 删除。"""

from __future__ import annotations

import logging
import shutil
from datetime import datetime, timezone
from pathlib import Path

from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session

from ..config import get_settings
from ..database import get_db
from ..models.task import Task, STATUS_UPLOADING, STATUS_DONE, STATUS_FAILED
from ..models.upload_session import UploadSession
from ..schemas.task import TaskListResponse, TaskResponse

logger = logging.getLogger("audio2text.tasks")
router = APIRouter(prefix="/api/tasks", tags=["task"])


def _to_resp(task: Task) -> TaskResponse:
    return TaskResponse(
        id=task.id,
        filename=task.filename,
        status=task.status,
        progress=task.progress,
        error=task.error,
        created_at=task.created_at,
        updated_at=task.updated_at,
    )


def _upload_to_resp(session: UploadSession) -> TaskResponse:
    """把上传中的 UploadSession 映射为虚拟 TaskResponse。

    is_upload=true 让前端走上传状态轮询而非任务轮询。
    progress = 已传分片数 / 总分片数 × 100（映射到 0-5 区间，与 extract 阶段衔接）。
    """
    uploaded = len(session.uploaded_chunks or [])
    total = session.total_chunks or 1
    # 上传进度映射到 0-4%（extract 从 5% 开始，留 1% 给 complete 拼接）
    progress = min(4.0, uploaded / total * 4.0)
    now = datetime.now(timezone.utc)
    return TaskResponse(
        id=0,  # 虚拟 id，前端用 upload_id 轮询
        filename=session.filename,
        status=STATUS_UPLOADING,
        progress=progress,
        error=None,
        created_at=session.created_at,
        updated_at=session.updated_at or now,
        size_bytes=session.size_bytes,
        is_upload=True,
        upload_id=session.upload_id,
    )


@router.get("", response_model=TaskListResponse, summary="任务列表")
def list_tasks(
    limit: int = Query(100, ge=1, le=500),
    offset: int = Query(0, ge=0),
    q: str = Query("", description="按文件名模糊搜索（大小写不敏感，匹配子串）"),
    db: Session = Depends(get_db),
) -> TaskListResponse:
    # 上传中的会话（pending 状态）也作为虚拟任务返回，让前端能看到上传进度
    upload_q = db.query(UploadSession).filter(UploadSession.status == "pending")
    if q.strip():
        like = f"%{q.strip()}%"
        upload_q = upload_q.filter(UploadSession.filename.ilike(like))
    uploads = upload_q.order_by(UploadSession.created_at.desc()).all()

    # 已创建的 Task（含 extracting/transcribing/.../done/failed）
    task_q = db.query(Task).order_by(Task.created_at.desc(), Task.id.desc())
    if q.strip():
        like = f"%{q.strip()}%"
        task_q = task_q.filter(Task.filename.ilike(like))
    total_tasks = task_q.count()
    tasks = task_q.offset(offset).limit(limit).all()

    # 合并：上传会话 + Task，按 created_at 倒序
    upload_resps = [_upload_to_resp(s) for s in uploads]
    task_resps = [_to_resp(t) for t in tasks]
    all_resps = upload_resps + task_resps
    all_resps.sort(key=lambda r: r.created_at, reverse=True)

    # 分页：offset/limit 作用于合并后列表（上传会话通常很少，主要影响首页前几条）
    paged = all_resps[offset:offset + limit]
    return TaskListResponse(tasks=paged, total=len(all_resps))


@router.get("/{task_id}", response_model=TaskResponse, summary="任务状态")
def get_task(task_id: int, db: Session = Depends(get_db)) -> TaskResponse:
    task = db.get(Task, task_id)
    if task is None:
        raise HTTPException(404, f"任务不存在：{task_id}")
    return _to_resp(task)


@router.get("/{task_id}/subtitle", summary="下载字幕")
def download_subtitle(
    task_id: int,
    type: str = Query("bilingual", pattern="^(bilingual|en|zh)$"),
    db: Session = Depends(get_db),
) -> FileResponse:
    task = db.get(Task, task_id)
    if task is None:
        raise HTTPException(404, f"任务不存在：{task_id}")
    if task.status != "done":
        raise HTTPException(409, f"任务尚未完成（当前状态：{task.status}）")

    rel = {
        "bilingual": task.bilingual_srt_path,
        "en": task.en_srt_path,
        "zh": task.zh_srt_path,
    }[type]
    if not rel:
        raise HTTPException(404, f"该类型字幕不存在：{type}")

    s = get_settings()
    path = s.output_dir() / rel
    if not path.is_file():
        raise HTTPException(404, f"字幕文件丢失：{path}")

    base = Path(task.filename).stem
    suffix = "" if type == "bilingual" else f".{type}"
    download_name = f"{base}{suffix}.srt"
    return FileResponse(
        path=path,
        media_type="application/x-subrip",
        filename=download_name,
    )


@router.delete("/{task_id}", summary="删除任务（仅允许已完成/失败）")
def delete_task(task_id: int, db: Session = Depends(get_db)) -> dict:
    """删除任务及其产物（字幕 / 中间音频 / 保留的原始视频）+ DB 记录。

    仅允许删除已完成（done）或失败（failed）的任务，进行中的任务不可删。
    """
    task = db.get(Task, task_id)
    if task is None:
        raise HTTPException(404, f"任务不存在：{task_id}")
    if task.status not in (STATUS_DONE, STATUS_FAILED):
        raise HTTPException(409, f"任务进行中，无法删除（当前状态：{task.status}）")

    s = get_settings()
    deleted: list[str] = []

    # 删字幕输出目录
    out_dir = s.output_dir() / f"task_{task.id}"
    if out_dir.is_dir():
        shutil.rmtree(out_dir, ignore_errors=True)
        deleted.append("outputs")

    # 删中间音频
    if task.wav_path:
        wav = Path(task.wav_path)
        if wav.is_file():
            wav.unlink(missing_ok=True)
            deleted.append("audio")

    # 删保留的原始视频
    if task.source_path:
        src = s.upload_dir() / task.source_path
        if src.is_file():
            src.unlink(missing_ok=True)
            deleted.append("video")

    # 删 DB 记录（先删关联的 UploadSession，再删 Task）
    db.query(UploadSession).filter(UploadSession.task_id == task.id).delete()
    db.delete(task)
    db.commit()

    logger.info("删除任务 %d（%s）：%s", task_id, task.filename, ", ".join(deleted) or "无产物")
    return {"status": "deleted", "task_id": task_id, "cleaned": deleted}