Agentic-OS/backend/app/api/tasks.py

"""Task creation, listing, approvals, and live event streaming."""
from __future__ import annotations

import asyncio
import json

from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.core.deps import require_readonly, require_support
from app.core.security import decode_access_token
from app.database import SessionLocal, get_db
from app.models.enums import ApprovalStatus, TaskStatus
from app.models.task import ApprovalRequest, Task
from app.models.user import User
from app.schemas import (
    ApprovalDecision,
    ApprovalOut,
    TaskCreate,
    TaskContinue,
    TaskDetail,
    TaskEscalateLlm,
    TaskEventOut,
    TaskOverviewOut,
    TaskOut,
    TaskPreviewIn,
    TaskPreviewOut,
)
from app.services.audit import record_audit
from app.services.events import enqueue_task, publish_event, subscribe_events
from app.inventory_catalog import catalog_by_key
from app.models.inventory import Vessel
from app.services.task_present import (
    task_to_detail,
    task_to_out,
    tasks_to_out_list,
    tasks_to_overview_list,
    vessel_name_map,
)
from app.services.task_runtime import CANCELLABLE_STATUSES, clear_task_cancel, request_task_cancel
from app.services.troubleshooting_rules import select_devices_for_issue

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

EVENT_REPLAY_LIMIT = 400

CHECK_LABELS = {
    "proxmox": ["VM and LXC status", "Guest configuration", "Host health"],
    "pfsense": ["System status", "Interfaces", "Gateways", "Firewall/NAT when requested"],
    "docker_vm": ["Container status", "Host resources"],
    "asterisk_geneseasx": ["Asterisk health", "SIP registrations", "Channels"],
    "asterisk_satbox": ["Asterisk health", "SIP registrations"],
    "fortigate": ["Firewall status"],
    "fortiswitch": ["Switch status"],
}


def _preview_device_dict(device) -> dict:
    key = device.catalog_key or device.device_type.value
    entry = catalog_by_key().get(device.catalog_key or "")
    return {
        "id": device.id,
        "name": device.name,
        "catalog_key": key,
        "device_type": device.device_type.value,
        "label": entry.label if entry else key.replace("_", " ").title(),
    }


def _checks_for_device(key: str | None) -> list[str]:
    return CHECK_LABELS.get(key or "", ["Health and status checks"])


@router.get("", response_model=list[TaskOut])
async def list_tasks(
    limit: int = 100,
    offset: int = 0,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_readonly),
):
    limit = max(1, min(limit, 500))
    offset = max(0, offset)
    res = await db.execute(
        select(Task).order_by(Task.id.desc()).limit(limit).offset(offset)
    )
    return await tasks_to_out_list(db, list(res.scalars().all()))


@router.get("/overview", response_model=list[TaskOverviewOut])
async def task_overview(
    limit: int = 50,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_readonly),
):
    """Compact task data for dashboard cards. No raw diagnostic payloads."""
    limit = max(1, min(limit, 200))
    res = await db.execute(select(Task).order_by(Task.id.desc()).limit(limit))
    return await tasks_to_overview_list(db, list(res.scalars().all()))


@router.post("/preview", response_model=TaskPreviewOut)
async def preview_task_scope(
    body: TaskPreviewIn,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_readonly),
):
    """Preview the devices/check groups that task creation will select."""
    res = await db.execute(
        select(Vessel).options(selectinload(Vessel.devices)).where(Vessel.id == body.vessel_id)
    )
    vessel = res.scalar_one_or_none()
    if not vessel:
        raise HTTPException(status_code=404, detail="Vessel not found")
    device_dicts = [_preview_device_dict(d) for d in vessel.devices]
    selected, rule = select_devices_for_issue(device_dicts, body.issue, title=body.title)
    warning = None
    if not selected:
        warning = "No onboard device matched this issue. Mention a device type or ask for a full health check."
    return TaskPreviewOut(
        vessel_id=vessel.id,
        vessel_name=vessel.name,
        devices=[
            {
                "catalog_key": d.get("catalog_key") or d.get("device_type") or "device",
                "name": d.get("name") or d.get("catalog_key") or "device",
                "label": d.get("label") or (d.get("catalog_key") or "device").replace("_", " ").title(),
                "checks": _checks_for_device(d.get("catalog_key")),
            }
            for d in selected
        ],
        rule_id=rule.id if rule else None,
        rule_name=rule.name if rule else None,
        severity=rule.severity if rule else None,
        hints=rule.hints if rule else [],
        warning=warning,
    )


@router.post("", response_model=TaskOut, status_code=status.HTTP_201_CREATED)
async def create_task(
    body: TaskCreate,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_support),
):
    if not body.vessel_id:
        raise HTTPException(status_code=400, detail="vessel_id is required — pick a vessel; the agent selects devices from the issue")
    task = Task(
        title=body.title,
        issue=body.issue,
        details=body.details,
        device_id=None,
        vessel_id=body.vessel_id,
        created_by_id=user.id,
        status=TaskStatus.queued,
    )
    db.add(task)
    await db.commit()
    await db.refresh(task)
    await enqueue_task(task.id)
    await publish_event(task.id, "log", "Task queued.", persist=True)
    await record_audit(db, user, "task.create", "task", task.id, {"title": task.title})
    vmap = await vessel_name_map(db, {task.vessel_id} if task.vessel_id else set())
    return task_to_out(task, vessel_name=vmap.get(task.vessel_id) if task.vessel_id else None)


@router.post("/{task_id}/cancel", response_model=TaskOut)
async def cancel_task(
    task_id: int,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_support),
):
    """Stop a queued or running task. In-flight work is interrupted at the next checkpoint."""
    res = await db.execute(select(Task).where(Task.id == task_id))
    task = res.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    if task.status.value not in CANCELLABLE_STATUSES:
        raise HTTPException(
            status_code=400,
            detail=f"Cannot stop a task with status '{task.status.value}'",
        )

    await request_task_cancel(task_id)
    task.status = TaskStatus.cancelled
    task.error = None
    await db.commit()
    await db.refresh(task)
    await publish_event(
        task_id,
        "cancelled",
        f"Task stopped by {user.email}.",
        {"requested_by": user.email},
        persist=True,
    )
    await record_audit(db, user, "task.cancel", "task", task.id)
    vmap = await vessel_name_map(db, {task.vessel_id} if task.vessel_id else set())
    return task_to_out(task, vessel_name=vmap.get(task.vessel_id) if task.vessel_id else None)


@router.post("/{task_id}/retry", response_model=TaskOut)
async def retry_task(
    task_id: int,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_support),
):
    res = await db.execute(select(Task).where(Task.id == task_id))
    task = res.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    if task.status in (TaskStatus.running, TaskStatus.queued):
        raise HTTPException(status_code=400, detail="Task is already running")
    await clear_task_cancel(task_id)
    task.status = TaskStatus.queued
    task.error = None
    task.report = None
    task.details = {}
    await db.commit()
    await db.refresh(task)
    await enqueue_task(task.id)
    await publish_event(task_id, "log", "Task re-queued from scratch (run history cleared).", persist=True)
    await record_audit(db, user, "task.retry", "task", task.id)
    return task


@router.post("/{task_id}/continue", response_model=TaskOut)
async def continue_task(
    task_id: int,
    body: TaskContinue,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_support),
):
    """Run another investigation pass on the same task with optional follow-up instructions."""
    res = await db.execute(select(Task).where(Task.id == task_id))
    task = res.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    if task.status in (TaskStatus.running, TaskStatus.queued):
        raise HTTPException(status_code=400, detail="Task is already running")
    if not body.follow_up.strip():
        raise HTTPException(status_code=400, detail="follow_up is required — describe what to check next")

    await clear_task_cancel(task_id)
    details = dict(task.details or {})
    details["_continue"] = {
        "follow_up": body.follow_up.strip(),
        "prior_report": task.report,
    }
    task.details = details
    task.status = TaskStatus.queued
    task.error = None
    await db.commit()
    await db.refresh(task)
    await enqueue_task(task.id)
    await publish_event(
        task_id,
        "log",
        f"Task queued for follow-up run: {body.follow_up.strip()[:200]}",
        persist=True,
    )
    await record_audit(db, user, "task.continue", "task", task.id, {"follow_up": body.follow_up[:500]})
    return task


@router.post("/{task_id}/escalate-llm", response_model=TaskOut)
async def escalate_task_llm(
    task_id: int,
    body: TaskEscalateLlm,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_support),
):
    """Request cloud API for reasoning — skips remaining local LLM tiers when possible."""
    from app.services.task_runtime import set_task_min_tier

    res = await db.execute(select(Task).where(Task.id == task_id))
    task = res.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    if task.status != TaskStatus.running:
        raise HTTPException(status_code=400, detail="Task is not running")

    await set_task_min_tier(task_id, body.tier)
    label = "Economy" if body.tier == "economy" else "Premium"
    await publish_event(
        task_id,
        "log",
        f"{label} cloud API requested by {user.email} — will skip local LLM for remaining reasoning steps.",
        {"min_tier": body.tier, "requested_by": user.email},
        persist=True,
    )
    await record_audit(db, user, "task.escalate_llm", "task", task.id, {"tier": body.tier})
    await db.refresh(task)
    return task


@router.get("/{task_id}", response_model=TaskDetail)
async def get_task(
    task_id: int, db: AsyncSession = Depends(get_db), _: User = Depends(require_readonly)
):
    res = await db.execute(
        select(Task).options(selectinload(Task.events)).where(Task.id == task_id)
    )
    task = res.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    vmap = await vessel_name_map(db, {task.vessel_id} if task.vessel_id else set())
    return task_to_detail(task, vessel_name=vmap.get(task.vessel_id) if task.vessel_id else None)


@router.get("/{task_id}/events", response_model=list[TaskEventOut])
async def get_task_events(
    task_id: int, db: AsyncSession = Depends(get_db), _: User = Depends(require_readonly)
):
    res = await db.execute(
        select(Task).options(selectinload(Task.events)).where(Task.id == task_id)
    )
    task = res.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    return task.events


@router.get("/{task_id}/approvals", response_model=list[ApprovalOut])
async def get_approvals(
    task_id: int, db: AsyncSession = Depends(get_db), _: User = Depends(require_readonly)
):
    res = await db.execute(select(ApprovalRequest).where(ApprovalRequest.task_id == task_id))
    return res.scalars().all()


@router.post("/{task_id}/approvals/{approval_id}", response_model=ApprovalOut)
async def decide_approval(
    task_id: int,
    approval_id: int,
    body: ApprovalDecision,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_support),
):
    res = await db.execute(
        select(ApprovalRequest).where(
            ApprovalRequest.id == approval_id, ApprovalRequest.task_id == task_id
        )
    )
    approval = res.scalar_one_or_none()
    if not approval:
        raise HTTPException(status_code=404, detail="Approval not found")
    from datetime import datetime, timezone

    approval.status = ApprovalStatus.approved if body.approve else ApprovalStatus.rejected
    approval.decided_by_id = user.id
    approval.decided_at = datetime.now(timezone.utc)

    if body.approve and approval.tool_name == "mcp_git_push" and approval.tool_args:
        from app.services import mcp_workspace

        args = approval.tool_args
        push = await mcp_workspace.commit_and_push(
            args.get("server", ""),
            args.get("commit_message", "agentic-os: MCP update"),
            args.get("files") or [],
        )
        if not push.get("ok"):
            if push.get("runtime_only"):
                raise HTTPException(
                    status_code=400,
                    detail=(
                        push.get("error")
                        or "Changes exist only in the running MCP — no git repo to push."
                    ),
                )
            raise HTTPException(
                status_code=502,
                detail=push.get("error", "Git push failed"),
            )
        note = push.get("note") or ""
        msg = f"MCP pushed to Gitea: {args.get('server')}"
        if push.get("commit"):
            msg += f" ({push.get('commit')})"
        elif note:
            msg += f" — {note}"
        await publish_event(
            task_id,
            "mcp_dev",
            msg,
            push,
        )

    await db.commit()
    await db.refresh(approval)
    await record_audit(
        db, user, "approval.decide", "approval", approval.id, {"approve": body.approve}
    )
    await publish_event(
        task_id,
        "approval",
        f"Approval {approval.id} {'approved' if body.approve else 'rejected'} by {user.email}",
        {"approval_id": approval.id, "approve": body.approve},
    )

    # If no more pending approvals, settle task status
    pend = await db.execute(
        select(ApprovalRequest).where(
            ApprovalRequest.task_id == task_id, ApprovalRequest.status == ApprovalStatus.pending
        )
    )
    if pend.scalars().first() is None:
        tres = await db.execute(select(Task).where(Task.id == task_id))
        task = tres.scalar_one_or_none()
        if task and task.status == TaskStatus.waiting_approval:
            task.status = TaskStatus.succeeded
            await db.commit()
    return approval


@router.websocket("/{task_id}/stream")
async def stream_task(websocket: WebSocket, task_id: int):
    """Live event stream. Authenticate via ?token=<jwt> query param."""
    token = websocket.query_params.get("token")
    payload = decode_access_token(token) if token else None
    if not payload:
        await websocket.close(code=4401)
        return
    await websocket.accept()

    # Replay the most recent persisted events first (bounded), oldest-first.
    async with SessionLocal() as db:
        from app.models.task import TaskEvent

        res = await db.execute(
            select(TaskEvent)
            .where(TaskEvent.task_id == task_id)
            .order_by(TaskEvent.id.desc())
            .limit(EVENT_REPLAY_LIMIT)
        )
        recent = list(reversed(res.scalars().all()))
        for ev in recent:
            await websocket.send_text(
                json.dumps(
                    {
                        "task_id": task_id,
                        "kind": ev.kind,
                        "message": ev.message,
                        "payload": ev.payload,
                        "ts": ev.created_at.isoformat(),
                        "replay": True,
                    }
                )
            )

    try:
        async for event in subscribe_events(task_id):
            await websocket.send_text(json.dumps(event))
    except WebSocketDisconnect:
        pass
    except asyncio.CancelledError:
        pass