Model/backend/app/tasks/router.py

from fastapi import APIRouter, Depends, HTTPException
from uuid import UUID
import json  # ← REQUIRED for json.loads

from backend.app.dependencies import validate_token
from backend.app.tasks.schema import (
    CreateTaskRequest,
    UpdateTaskStatusRequest,
    CreateSubTaskRequest,
    UpdateSubTaskStatusRequest,
    FinalizeSubTaskRequest,
    TaskSqsTriggerRequest
)

# Correct location of interfaces
from backend.app.db.functions.tasks.Tasks import TasksInterface, SubTaskInterface

from backend.app.db.connection import get_db_session
from backend.app.db.models.tasks import Task, SubTask
from sqlmodel import select


router = APIRouter(
    prefix="/tasks",
    tags=["tasks"],
    dependencies=[Depends(validate_token)],
)


# ============================================================
# Create Task
# ============================================================
@router.post("/", summary="Create a new task and its first subtask")
async def create_task(req: CreateTaskRequest):
    tasks = TasksInterface()
    task_id, subtask_id = tasks.create_task(
        task_source=req.task_source,
        service=req.service,
        inputs=req.inputs,
    )
    return {"task_id": task_id, "subtask_id": subtask_id}


# ============================================================
# Get Task + Subtasks
# ============================================================
@router.get("/{task_id}", summary="Get a task and its subtasks")
async def get_task(task_id: UUID):
    with get_db_session() as session:
        task = session.get(Task, task_id)
        if not task:
            raise HTTPException(status_code=404, detail="Task not found")

        subtasks = session.exec(
            select(SubTask).where(SubTask.taskId == task_id)
        ).all()

        formatted = []
        for st in subtasks:
            formatted.append({
                **st.dict(),
                "inputs": json.loads(st.inputs) if st.inputs else None,
                "outputs": json.loads(st.outputs) if st.outputs else None,
                "cloud_logs_url": st.cloudLogsURL,
            })

        return {
            "task": task,
            "subtasks": formatted,
        }


# ============================================================
# Update Task Status
# ============================================================
@router.put("/{task_id}/status", summary="Update a task's status")
async def update_task_status(task_id: UUID, req: UpdateTaskStatusRequest):
    tasks = TasksInterface()
    try:
        updated = tasks.update_task_status(task_id, req.status)
        return {"task_id": updated.id, "status": updated.status}
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


# ============================================================
# Create Additional Subtask
# ============================================================
@router.post("/{task_id}/subtasks", summary="Create a new subtask under a task")
async def create_subtask(task_id: UUID, req: CreateSubTaskRequest):
    subtasks = SubTaskInterface()
    try:
        st = subtasks.create_subtask(task_id, req.inputs)
        return {"subtask_id": st.id, "task_id": task_id, "status": st.status}
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


# ============================================================
# Update Subtask Status
# ============================================================
@router.put("/subtask/{subtask_id}/status", summary="Update a subtask's status")
async def update_subtask_status(subtask_id: UUID, req: UpdateSubTaskStatusRequest):
    subtasks = SubTaskInterface()
    try:
        st = subtasks.update_subtask_status(subtask_id, req.status)
        return {"subtask_id": st.id, "status": st.status}
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


# ===
# Sub task is complete
@router.post("/subtask/{subtask_id}/finalize", summary="Finalize a subtask with status, outputs, logs")
async def finalize_subtask(subtask_id: UUID, req: FinalizeSubTaskRequest):
    subtasks = SubTaskInterface()

    try:
        st = subtasks.finalize_subtask(
            subtask_id=subtask_id,
            status=req.status,
            outputs=req.outputs,
            cloud_logs_url=req.cloud_logs_url
        )

        return {
            "subtask_id": st.id,
            "status": st.status,
            "outputs": req.outputs,
            "cloud_logs_url": req.cloud_logs_url,
        }

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


# for testing:

import boto3
import json
from backend.app.tasks.schema import TaskSqsTriggerRequest
from backend.app.db.functions.tasks.Tasks import TasksInterface, SubTaskInterface
from backend.app.config import get_settings

sqs = boto3.client("sqs")

@router.post("/trigger", summary="Create task + subtask and publish to SQS", status_code=202)
async def trigger_task(req: TaskSqsTriggerRequest):
    """
    Creates a Task + SubTask, then pushes the SubTask into SQS so a Lambda can process it.
    If inputs are empty, automatically replaced with {}.
    """

    settings = get_settings()

    tasks = TasksInterface()

    # ---- Normalize empty inputs ----
    inputs = req.inputs or {}   # ensures {} even if null

    # ---- 1. Create Task + SubTask ----
    task_id, subtask_id = tasks.create_task(
        task_source=req.task_source,
        service=req.service,
        inputs=inputs,
    )

    # ---- 2. Prepare SQS payload ----
    sqs_payload = {
        "subtask_id": str(subtask_id),
        "params": inputs,
    }

    try:
        response = sqs.send_message(
            QueueUrl=f"https://sqs.{settings.AWS_REGION}.amazonaws.com/"
                     f"{settings.AWS_ACCOUNT_ID}/lambda-example-queue",
            MessageBody=json.dumps(sqs_payload)
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"SQS error: {e}")

    return {
        "message": "Task triggered",
        "task_id": task_id,
        "subtask_id": subtask_id,
        "sqs_message_id": response.get("MessageId"),
        "inputs_sent": inputs,
    }