yu-zhi-ran/platform/backend/app/api/tasks.py

import uuid
from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.orm import Session
from typing import List, Optional

from ..database import get_db
from ..models import ContentTask, Topic
from ..schemas import ContentTaskCreate, ContentTaskResponse
from .auth import get_current_user, org_filter

router = APIRouter(prefix="/api/tasks", tags=["tasks"])


@router.get("", response_model=List[ContentTaskResponse])
def list_tasks(
    status: Optional[str] = None,
    topic_id: Optional[str] = None,
    stage: Optional[str] = None,
    limit: int = Query(50, le=200),
    offset: int = Query(0, ge=0),
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    query = db.query(ContentTask, Topic.title.label("topic_title")).join(Topic, ContentTask.topic_id == Topic.id, isouter=True)
    of = org_filter(current_user, Topic)
    if of is not True:
        query = query.filter((ContentTask.topic_id.is_(None)) | (Topic.org_id == current_user.org_id))
    if status:
        query = query.filter(ContentTask.status == status)
    if topic_id:
        query = query.filter(ContentTask.topic_id == topic_id)
    if stage:
        query = query.filter(ContentTask.stage == stage)
    rows = query.order_by(ContentTask.created_at.desc()).offset(offset).limit(limit).all()
    result = []
    for r in rows:
        task = r.ContentTask
        task.topic_title = r.topic_title
        result.append(task)
    return result


@router.get("/active")
def get_active_tasks(
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    q = db.query(ContentTask).join(Topic, ContentTask.topic_id == Topic.id, isouter=True)
    of = org_filter(current_user, Topic)
    if of is not True:
        q = q.filter((ContentTask.topic_id.is_(None)) | (Topic.org_id == current_user.org_id))
    return q.filter(
        ContentTask.status == "running"
    ).order_by(ContentTask.started_at.desc()).all()


@router.post("", response_model=ContentTaskResponse)
def create_task(
    data: ContentTaskCreate,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    if data.topic_id:
        topic = db.query(Topic).filter(Topic.id == data.topic_id).first()
        if not topic:
            raise HTTPException(status_code=404, detail="选题不存在")
        if current_user.role != "admin" and topic.org_id != current_user.org_id:
            raise HTTPException(status_code=404, detail="选题不存在")

    task_id = f"task_{uuid.uuid4().hex[:16]}"

    task = ContentTask(
        task_id=task_id,
        topic_id=data.topic_id,
        stage=data.stage,
        status="pending",
        created_by=data.created_by or current_user.username
    )
    db.add(task)
    db.commit()
    db.refresh(task)
    return task


@router.get("/{task_id}", response_model=ContentTaskResponse)
def get_task(
    task_id: str,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    task = db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")
    return task


@router.put("/{task_id}/start")
def start_task(
    task_id: str,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    task = db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")

    from datetime import datetime, timezone
    task.status = "running"
    task.started_at = datetime.now(timezone.utc)
    task.message = "任务已启动"
    task.progress = 0
    db.commit()
    db.refresh(task)
    return task


@router.put("/{task_id}/progress")
def update_progress(
    task_id: str,
    progress: int = Query(..., ge=0, le=100),
    message: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    task = db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")

    task.progress = progress
    if message:
        task.message = message
    db.commit()
    db.refresh(task)
    return task


@router.put("/{task_id}/complete")
def complete_task(
    task_id: str,
    result_data: dict = None,
    message: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    task = db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")

    from datetime import datetime, timezone
    finished = datetime.now(timezone.utc)
    task.status = "completed"
    task.finished_at = finished
    task.progress = 100
    if message:
        task.message = message
    if result_data:
        task.result_data = result_data
    if task.started_at:
        task.duration = int((finished - task.started_at).total_seconds())
    db.commit()
    db.refresh(task)
    return task


@router.put("/{task_id}/fail")
def fail_task(
    task_id: str,
    error_msg: str = Query(...),
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    task = db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")

    from datetime import datetime, timezone
    finished = datetime.now(timezone.utc)
    task.status = "failed"
    task.finished_at = finished
    task.error_msg = error_msg
    if task.started_at:
        task.duration = int((finished - task.started_at).total_seconds())
    db.commit()
    db.refresh(task)
    return task


@router.delete("/{task_id}")
def cancel_task(
    task_id: str,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    task = db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")

    task.status = "cancelled"
    db.commit()
    return {"ok": True}


@router.post("/run-creator")
def run_creator_task(
    topic_id: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user=Depends(get_current_user)
):
    import threading
    from datetime import datetime, timezone

    now = datetime.now(timezone.utc)
    task_id = f"task_{uuid.uuid4().hex[:16]}"

    task = ContentTask(
        task_id=task_id,
        topic_id=topic_id,
        stage="creator",
        status="running",
        started_at=now,
        created_by=current_user.username
    )
    db.add(task)
    db.commit()
    db.refresh(task)

    from ..core.generator import run_creator

    def _run():
        from ..database import SessionLocal
        from datetime import datetime, timezone
        new_db = SessionLocal()
        try:
            new_task = new_db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
            if new_task:
                new_task.message = "创作脚本运行中..."
                new_task.progress = 30
                new_db.commit()
            result = run_creator(topic_id)
            new_task = new_db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
            if new_task:
                finished = datetime.now(timezone.utc)
                new_task.status = "completed"
                new_task.finished_at = finished
                new_task.progress = 100
                new_task.message = "创作完成"
                new_task.result_data = result or {}
                if new_task.started_at:
                    new_task.duration = int((finished - new_task.started_at).total_seconds())
                new_db.commit()
        except Exception as e:
            new_task = new_db.query(ContentTask).filter(ContentTask.task_id == task_id).first()
            if new_task:
                finished = datetime.now(timezone.utc)
                new_task.status = "failed"
                new_task.finished_at = finished
                new_task.error_msg = str(e)
                if new_task.started_at:
                    new_task.duration = int((finished - new_task.started_at).total_seconds())
                new_db.commit()
        finally:
            new_db.close()

    thread = threading.Thread(target=_run)
    thread.start()

    return task