api/tasks.py

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import Optional

from core.task_manager import create_tasks, complete_task, get_task_status, delete_completed_tasks
from core.state import tasks, get_task
from config import DATA_SOURCES

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

class GoalRequest(BaseModel):
    target_size: str
    language: Optional[str] = "python"
    data_source: Optional[str] = None
    total_steps: Optional[int] = 1000
    steps_per_task: Optional[int] = 100

class TaskResultRequest(BaseModel):
    worker_id: str
    token: str
    task_id: str
    result: dict
    loss: Optional[float] = None
    steps: Optional[int] = None

@router.post("/create")
async def create(data: GoalRequest):
    try:
        from config import config
        config.set("total_steps", data.total_steps or 1000)
        config.set("steps_per_task", data.steps_per_task or 100)
        result = create_tasks(
            target_size=data.target_size,
            language=data.language,
            data_source=data.data_source
        )
        return {"status": "created", **result}
    except ValueError as e:
        raise HTTPException(400, str(e))
    except Exception as e:
        print(f"❌ Create task error: {e}")
        raise HTTPException(500, str(e))

@router.post("/complete")
async def complete(data: TaskResultRequest):
    try:
        success = complete_task(
            task_id=data.task_id,
            worker_id=data.worker_id,
            loss=data.loss or 0.0,
            steps=data.steps or 0,
            result_data=data.result
        )
        if not success:
            raise HTTPException(404, "Task not found or not assigned")
        return {"status": "ok"}
    except HTTPException:
        raise
    except Exception as e:
        print(f"❌ Complete task error: {e}")
        raise HTTPException(500, str(e))

@router.get("/status")
async def status():
    try:
        return get_task_status()
    except Exception as e:
        print(f"❌ Status error: {e}")
        raise HTTPException(500, str(e))

@router.get("/list")
async def list_tasks():
    try:
        return {"tasks": tasks, "total": len(tasks)}
    except Exception as e:
        print(f"❌ List tasks error: {e}")
        raise HTTPException(500, str(e))

@router.post("/delete-completed")
async def delete_completed():
    try:
        count = delete_completed_tasks()
        return {"deleted": count, "remaining": len(tasks)}
    except Exception as e:
        print(f"❌ Delete completed error: {e}")
        raise HTTPException(500, str(e))