CapacityReport/app/api/routers/tasks.py

from datetime import datetime
from pathlib import Path
from threading import Thread
from typing import Any

from fastapi import APIRouter, Body, HTTPException

from app import state
from app.api.routers.task_runtime import (
    apply_history_retention_safely,
    log_license_check,
    set_task_stage,
)
from app.config import AppConfig
from app.processor import DataProcessor, ProcessLogger
from app.services.license import LicenseError, check_processing_allowed


router = APIRouter(tags=["tasks"])


@router.get("/api/task/status")
async def get_global_task_status():
    if state.global_task_lock["locked"]:
        task_id = state.global_task_lock["task_id"]
        if task_id and _task_finished(task_id):
            state.reset_task_lock()
            return {"has_active": False}

        return {
            "has_active": True,
            "task_id": task_id,
            "stage": state.global_task_lock["stage"],
            "started_at": state.global_task_lock["started_at"],
            "logs": [],
        }

    active_tasks = {
        task_id: task
        for task_id, task in state.processing_tasks.items()
        if task.get("status") == "processing"
    }
    if active_tasks:
        task_id = next(iter(active_tasks))
        return {
            "has_active": True,
            "task_id": task_id,
            "stage": active_tasks[task_id].get("stage", "processing"),
            "logs": active_tasks[task_id].get("logs", []),
        }

    return {"has_active": False}


@router.post("/api/task/lock")
async def lock_task(task_id: str = Body(..., embed=True)):
    if state.global_task_lock["locked"]:
        raise HTTPException(status_code=409, detail="已有任务在运行")

    state.global_task_lock.update(
        {
            "locked": True,
            "task_id": task_id,
            "stage": "uploading",
            "started_at": datetime.now().isoformat(),
        }
    )
    return {"success": True, "message": "任务已锁定"}


@router.post("/api/task/unlock")
async def unlock_task(task_id: str | None = Body(None, embed=True)):
    if task_id and state.global_task_lock["task_id"] != task_id:
        raise HTTPException(status_code=403, detail="无权解锁此任务")

    state.reset_task_lock()
    return {"success": True, "message": "任务已解锁"}


@router.get("/api/process/active")
async def get_active_task():
    return await get_global_task_status()


@router.post("/api/process/start")
async def start_processing(task_id: str = Body(..., embed=True)):
    record = state.history_manager.get(task_id)
    if not record:
        raise HTTPException(status_code=404, detail="任务不存在")
    if record.status == "processing":
        raise HTTPException(status_code=400, detail="任务正在处理中")

    work_dir = Path(record.work_dir)
    if not work_dir.exists():
        raise HTTPException(status_code=400, detail="工作目录不存在")

    logs: list[str] = []
    current_stage = "processing"

    def log_callback(message: str) -> None:
        logs.append(message)
        set_task_stage(task_id, current_stage, logs)

    def stage_callback(stage: str) -> None:
        nonlocal current_stage
        current_stage = stage
        set_task_stage(task_id, current_stage, logs)

    logger = ProcessLogger(
        log_file=work_dir / "log.txt",
        callback=log_callback,
        stage_callback=stage_callback,
    )
    app_config = state.current_config()
    state.history_manager.update(task_id, status="processing")
    state.processing_tasks[task_id] = {"logs": [], "status": "processing", "stage": current_stage}
    state.global_task_lock.update(
        {
            "locked": True,
            "task_id": task_id,
            "stage": "processing",
            "started_at": datetime.now().isoformat(),
        }
    )

    thread = Thread(target=_run_processing, args=(task_id, work_dir, logger, app_config), daemon=True)
    thread.start()

    return {"success": True, "message": "处理任务已启动", "task_id": task_id}


@router.post("/api/process/status")
async def get_processing_status(task_id: str = Body(..., embed=True)):
    if task_id in state.processing_tasks:
        task_info = state.processing_tasks[task_id]
        logs = task_info.get("logs") or state.history_manager.get_logs(task_id)
        return {
            "task_id": task_id,
            "status": task_info["status"],
            "stage": task_info.get("stage"),
            "logs": logs,
            "error": task_info.get("error"),
            "error_detail": task_info.get("error_detail"),
        }

    record = state.history_manager.get(task_id)
    if not record:
        raise HTTPException(status_code=404, detail="任务不存在")

    return {
        "task_id": task_id,
        "status": record.status,
        "stage": record.status,
        "logs": state.history_manager.get_logs(task_id),
        "elapsed_time": record.elapsed_time,
        "error": record.error,
    }


def _task_finished(task_id: str) -> bool:
    record = state.history_manager.get(task_id)
    if record and record.status in {"completed", "failed"}:
        return True

    task_info: dict[str, Any] | None = state.processing_tasks.get(task_id)
    return bool(task_info and task_info.get("status") in {"completed", "failed"})


def _run_processing(task_id: str, work_dir: Path, logger: ProcessLogger, app_config: AppConfig) -> None:
    try:
        if app_config.warehouse_type == "metrix":
            import time

            from app.services.pipeline import RESULT_TABLES, run_import_and_report

            started = time.time()
            run_import_and_report(work_dir, app_config, logger)
            status, error, elapsed, result_tables = "completed", None, round(time.time() - started, 2), RESULT_TABLES
        else:
            logger.set_stage("license")
            log_license_check(logger, check_processing_allowed(work_dir))
            processor = DataProcessor(app_config, work_dir, logger)
            result = processor.process()
            status = "completed" if result.get("success") else "failed"
            error = result.get("error")
            elapsed = result.get("elapsed_time", 0)
            result_tables = ["4G_结果表", "5G_结果表"]
        state.history_manager.update(
            task_id,
            status=status,
            elapsed_time=elapsed,
            error=error,
            result_tables=result_tables,
        )
        state.processing_tasks[task_id] = {
            "logs": state.history_manager.get_logs(task_id),
            "status": status,
            "stage": status,
            "error": error,
        }
    except Exception as exc:
        error_detail = exc.to_detail() if isinstance(exc, LicenseError) else None
        logger.error(str(exc))
        state.history_manager.update(task_id, status="failed", error=str(exc))
        state.processing_tasks[task_id] = {
            "logs": state.history_manager.get_logs(task_id),
            "status": "failed",
            "stage": "failed",
            "error": str(exc),
            "error_detail": error_detail,
        }
    finally:
        apply_history_retention_safely()
        state.reset_task_lock()