utp_service/backend/app/routers/tasks.py

from typing import Optional

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

from app.dependencies import get_current_user, get_db
from app.models.device import Device
from app.models.object import Object
from app.models.rack import Rack
from app.models.task import Task, TaskResult
from app.models.zone import Zone
from app.models.user import User
from app.plugins.registry import action_registry
from app.schemas.task import ActionInfo, TaskCreate, TaskDetail, TaskPreviewResponse, TaskRead
from app.services import task as task_service

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


async def _resolve_labels(db: AsyncSession, tasks: list[Task]) -> dict[str, str]:
    """Batch-resolve target_scope to human-readable labels."""
    obj_ids: set[int] = set()
    dev_ids: set[int] = set()
    rack_ids: set[int] = set()
    zone_ids: set[int] = set()

    for t in tasks:
        scope = t.target_scope
        kind = scope.get("type")
        if kind == "object":
            obj_ids.add(scope["id"])
        elif kind == "device":
            dev_ids.add(scope["id"])
        elif kind == "device_list":
            for did in scope.get("ids", []):
                dev_ids.add(did)
        elif kind == "rack":
            rack_ids.add(scope["id"])
        elif kind == "zone":
            zone_ids.add(scope["id"])
        elif kind in ("category", "role") and (oid := scope.get("object_id")):
            obj_ids.add(oid)

    objs: dict[int, str] = {}
    devs: dict[int, str] = {}
    rack_names: dict[int, str] = {}
    zone_names: dict[int, str] = {}

    if obj_ids:
        rows = await db.execute(select(Object.id, Object.name).where(Object.id.in_(obj_ids)))
        objs = {row[0]: row[1] for row in rows.all()}

    if dev_ids:
        rows = await db.execute(
            select(Device.id, Device.hostname, Device.ip).where(Device.id.in_(dev_ids))
        )
        devs = {row[0]: (row[1] or row[2]) for row in rows.all()}

    if rack_ids:
        rows = await db.execute(select(Rack.id, Rack.name).where(Rack.id.in_(rack_ids)))
        rack_names = {row[0]: row[1] for row in rows.all()}

    if zone_ids:
        rows = await db.execute(select(Zone.id, Zone.name).where(Zone.id.in_(zone_ids)))
        zone_names = {row[0]: row[1] for row in rows.all()}

    labels: dict[str, str] = {}
    for t in tasks:
        scope = t.target_scope
        kind = scope.get("type", "")
        if kind == "object":
            labels[t.id] = objs.get(scope.get("id"), f"Объект #{scope.get('id')}")
        elif kind == "device":
            labels[t.id] = devs.get(scope.get("id"), f"Устройство #{scope.get('id')}")
        elif kind == "rack":
            _rid = scope.get('id')
            labels[t.id] = f"Стойка: {rack_names.get(_rid, f'#{_rid}')}"
        elif kind == "zone":
            _zid = scope.get('id')
            labels[t.id] = f"Зона: {zone_names.get(_zid, f'#{_zid}')}"
        elif kind == "device_list":
            ids = scope.get("ids", [])
            names = [devs.get(did, f"#{did}") for did in ids[:3]]
            suffix = f" и ещё {len(ids) - 3}" if len(ids) > 3 else ""
            labels[t.id] = ", ".join(names) + suffix
        elif kind in ("category", "role"):
            val = scope.get("value", "?")
            if oid := scope.get("object_id"):
                labels[t.id] = f"{val} @ {objs.get(oid, f'#{oid}')}"
            else:
                labels[t.id] = f"{val} (все объекты)"
        else:
            labels[t.id] = "—"

    return labels


@router.get("/actions", response_model=list[ActionInfo])
async def list_actions(_: User = Depends(get_current_user)):
    return [
        ActionInfo(
            name=a.name,
            display_name=a.display_name,
            compatible_categories=a.compatible_categories,
            params_schema=a.params_schema,
        )
        for a in action_registry.all()
    ]


@router.post("/preview", response_model=TaskPreviewResponse)
async def preview_task(
    body: TaskCreate,
    _: User = Depends(get_current_user),
):
    """Resolve scope and return device count breakdown before creating a task."""
    action = action_registry.get(body.type)
    if action is None:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Unknown action: {body.type}")

    targets = await task_service._resolve_scope(body.target_scope)

    by_category: dict[str, int] = {}
    for d, _ in targets:
        by_category[d.category] = by_category.get(d.category, 0) + 1

    if action.compatible_categories:
        compatible = [(d, o) for d, o in targets if action.is_compatible(d.category)]
        skipped = [(d, o) for d, o in targets if not action.is_compatible(d.category)]
    else:
        compatible = targets
        skipped = []

    compatible_by_category: dict[str, int] = {}
    for d, _ in compatible:
        compatible_by_category[d.category] = compatible_by_category.get(d.category, 0) + 1

    skipped_categories = sorted({d.category for d, _ in skipped})

    return TaskPreviewResponse(
        total=len(targets),
        compatible=len(compatible),
        skipped=len(skipped),
        by_category=by_category,
        compatible_by_category=compatible_by_category,
        skipped_categories=skipped_categories,
    )


@router.post("", response_model=TaskRead, status_code=status.HTTP_202_ACCEPTED)
async def create_task(
    body: TaskCreate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    try:
        task = await task_service.create_task(
            action_type=body.type,
            target_scope=body.target_scope,
            params=body.params,
            user_id=current_user.id,
        )
    except ValueError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
    labels = await _resolve_labels(db, [task])
    result = TaskRead.model_validate(task)
    result.target_label = labels.get(task.id)
    return result


@router.get("", response_model=list[TaskRead])
async def list_tasks(
    limit: int = 50,
    skip: int = 0,
    device_id: Optional[int] = Query(None),
    db: AsyncSession = Depends(get_db),
    _: User = Depends(get_current_user),
):
    query = select(Task).order_by(desc(Task.created_at))
    if device_id is not None:
        query = (
            query
            .join(TaskResult, Task.id == TaskResult.task_id)
            .where(TaskResult.device_id == device_id)
            .distinct()
        )
    rows = await db.execute(query.offset(skip).limit(limit))
    tasks = list(rows.scalars().all())
    labels = await _resolve_labels(db, tasks)
    result = []
    for t in tasks:
        r = TaskRead.model_validate(t)
        r.target_label = labels.get(t.id)
        result.append(r)
    return result


@router.get("/{task_id}", response_model=TaskDetail)
async def get_task(
    task_id: str,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(get_current_user),
):
    rows = await db.execute(
        select(Task)
        .options(selectinload(Task.events), selectinload(Task.device_results))
        .where(Task.id == task_id)
    )
    task = rows.scalar_one_or_none()
    if task is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Task not found")
    labels = await _resolve_labels(db, [task])
    result = TaskDetail.model_validate(task)
    result.target_label = labels.get(task.id)

    # Enrich device_results with IP, hostname, object, rack
    if result.device_results:
        device_ids = [r.device_id for r in result.device_results]
        dev_rows = await db.execute(
            select(Device.id, Device.ip, Device.hostname, Device.object_id, Device.rack_id)
            .where(Device.id.in_(device_ids))
        )
        dev_map = {
            row[0]: {"ip": row[1], "hostname": row[2], "object_id": row[3], "rack_id": row[4]}
            for row in dev_rows.all()
        }

        obj_ids = {v["object_id"] for v in dev_map.values() if v["object_id"]}
        rack_ids = {v["rack_id"] for v in dev_map.values() if v["rack_id"]}

        obj_names: dict[int, str] = {}
        rack_name_map: dict[int, str] = {}

        if obj_ids:
            obj_rows = await db.execute(
                select(Object.id, Object.name).where(Object.id.in_(obj_ids))
            )
            obj_names = {r[0]: r[1] for r in obj_rows.all()}

        if rack_ids:
            rack_rows = await db.execute(
                select(Rack.id, Rack.name).where(Rack.id.in_(rack_ids))
            )
            rack_name_map = {r[0]: r[1] for r in rack_rows.all()}

        for r in result.device_results:
            dev_info = dev_map.get(r.device_id, {})
            r.device_ip = dev_info.get("ip")
            r.device_hostname = dev_info.get("hostname")
            r.object_id = dev_info.get("object_id")
            r.rack_id = dev_info.get("rack_id")
            r.object_name = obj_names.get(dev_info.get("object_id")) if dev_info.get("object_id") else None
            r.rack_name = rack_name_map.get(dev_info.get("rack_id")) if dev_info.get("rack_id") else None

    return result


@router.post("/{task_id}/cancel", status_code=status.HTTP_204_NO_CONTENT)
async def cancel_task(
    task_id: str,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(get_current_user),
):
    task = await db.get(Task, task_id)
    if task is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Task not found")
    if task.status not in ("pending", "running"):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Cannot cancel task in status '{task.status}'",
        )
    task.cancel_requested = True