from fastapi import APIRouter, HTTPException

from app.schemas import GenerateRequest, HealthResponse, TaskResultResponse, TaskStatusResponse


router = APIRouter()


@router.post("/generate", response_model=TaskStatusResponse)
async def create_generate_task(req: GenerateRequest):
    task_manager = router.task_manager
    return await task_manager.create_task(req)


@router.get("/tasks/{task_id}", response_model=TaskStatusResponse)
def get_task_status(task_id: str):
    task_manager = router.task_manager
    try:
        return task_manager.get_status(task_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc


@router.get("/tasks/{task_id}/result", response_model=TaskResultResponse)
def get_task_result(task_id: str):
    task_manager = router.task_manager
    try:
        return task_manager.get_result(task_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc


@router.get("/health", response_model=HealthResponse)
def health_check():
    torch = router.torch
    ltx_backend = router.ltx_backend
    hunyuan_backend = router.hunyuan_backend

    cuda_ok = bool(torch.cuda.is_available())
    gpu_name = torch.cuda.get_device_name(0) if cuda_ok else None

    return HealthResponse(
        service_status="ok",
        cuda_available=cuda_ok,
        gpu_name=gpu_name,
        ltx_loaded=ltx_backend.is_loaded(),
        hunyuan_loaded=hunyuan_backend.is_loaded(),
    )