planet/backend/app/api/v1/tasks.py

from datetime import UTC, datetime
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import text

from app.db.session import get_db
from app.models.user import User
from app.core.security import get_current_user
from app.core.time import to_iso8601_utc
from app.services.collectors.registry import collector_registry


router = APIRouter()


@router.get("")
async def list_tasks(
    datasource_id: int = None,
    status: str = None,
    page: int = 1,
    page_size: int = 20,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    offset = (page - 1) * page_size
    query = """
        SELECT ct.id, ct.datasource_id, ds.name as datasource_name, ct.status,
               ct.started_at, ct.completed_at, ct.records_processed, ct.error_message,
               ct.phase, ct.phase_progress, ct.phase_message, ct.phase_current,
               ct.phase_total, ct.phase_unit, ct.total_records, ct.progress,
               ct.task_type, ct.source, ds.source as datasource_source
        FROM collection_tasks ct
        JOIN data_sources ds ON ct.datasource_id = ds.id
        WHERE 1=1
    """
    count_query = "SELECT COUNT(*) FROM collection_tasks ct WHERE 1=1"
    params = {}

    if datasource_id:
        query += " AND ct.datasource_id = :datasource_id"
        count_query += " AND ct.datasource_id = :datasource_id"
        params["datasource_id"] = datasource_id
    if status:
        statuses = [item.strip() for item in status.split(",") if item.strip()]
        if len(statuses) > 1:
            placeholders = ", ".join(f":status_{index}" for index, _item in enumerate(statuses))
            query += f" AND ct.status IN ({placeholders})"
            count_query += f" AND ct.status IN ({placeholders})"
            params.update({f"status_{index}": item for index, item in enumerate(statuses)})
        else:
            query += " AND ct.status = :status"
            count_query += " AND ct.status = :status"
            params["status"] = statuses[0] if statuses else status

    query += f" ORDER BY ct.created_at DESC LIMIT {page_size} OFFSET {offset}"

    result = await db.execute(text(query), params)
    tasks = result.fetchall()

    count_result = await db.execute(text(count_query), params)
    total = count_result.scalar()

    return {
        "total": total or 0,
        "page": page,
        "page_size": page_size,
        "data": [
            {
                "id": t[0],
                "datasource_id": t[1],
                "datasource_name": t[2],
                "status": t[3],
                "started_at": to_iso8601_utc(t[4]),
                "completed_at": to_iso8601_utc(t[5]),
                "records_processed": t[6],
                "error_message": t[7],
                "phase": t[8],
                "phase_progress": t[9],
                "phase_message": t[10],
                "phase_current": t[11],
                "phase_total": t[12],
                "phase_unit": t[13],
                "total_records": t[14],
                "progress": t[15],
                "task_type": t[16],
                "source": t[17] or t[18],
                "datasource_source": t[18],
            }
            for t in tasks
        ],
    }


@router.get("/{task_id}")
async def get_task(
    task_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(
        text("""
            SELECT ct.id, ct.datasource_id, ds.name as datasource_name, ct.status,
                   ct.started_at, ct.completed_at, ct.records_processed, ct.error_message,
                   ct.phase, ct.phase_progress, ct.phase_message, ct.phase_current,
                   ct.phase_total, ct.phase_unit, ct.total_records, ct.progress
            FROM collection_tasks ct
            JOIN data_sources ds ON ct.datasource_id = ds.id
            WHERE ct.id = :id
        """),
        {"id": task_id},
    )
    task = result.fetchone()

    if not task:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Task not found",
        )

    return {
        "id": task[0],
        "datasource_id": task[1],
        "datasource_name": task[2],
        "status": task[3],
        "started_at": to_iso8601_utc(task[4]),
        "completed_at": to_iso8601_utc(task[5]),
        "records_processed": task[6],
        "error_message": task[7],
        "phase": task[8],
        "phase_progress": task[9],
        "phase_message": task[10],
        "phase_current": task[11],
        "phase_total": task[12],
        "phase_unit": task[13],
        "total_records": task[14],
        "progress": task[15],
    }


@router.post("/datasources/{source_id}/trigger")
async def trigger_collection(
    source_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(
        text("SELECT id, name, collector_class FROM data_sources WHERE id = :id"),
        {"id": source_id},
    )
    datasource = result.fetchone()

    if not datasource:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Data source not found",
        )

    collector_class_name = datasource[2]
    collector_name = collector_class_name.lower().replace("collector", "")

    collector = collector_registry.get(collector_name)
    if not collector:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Collector {collector_name} not found",
        )

    result = await collector.run(db)

    await db.execute(
        text("""
            INSERT INTO collection_tasks (datasource_id, status, records_processed, error_message, started_at, completed_at, created_at)
            VALUES (:datasource_id, :status, :records_processed, :error_message, :started_at, :completed_at, NOW())
        """),
        {
            "datasource_id": source_id,
            "status": result.get("status", "unknown"),
            "records_processed": result.get("records_processed", 0),
            "error_message": result.get("error"),
            "started_at": datetime.now(UTC),
            "completed_at": datetime.now(UTC),
        },
    )

    return {
        "message": "Collection task executed",
        "result": result,
    }