"""
Runs once at app startup:
  - Marks tasks stuck in 'running' with stale heartbeat as 'stale'.
    These are tasks that were running when the app last crashed or restarted.
"""

from datetime import datetime, timedelta, timezone

import sqlalchemy as sa
from sqlalchemy import update

from app.models.task import Task
from app.services.database import AsyncSessionLocal

STALE_AFTER_MINUTES = 15


async def recover_stale_tasks() -> int:
    cutoff = datetime.now(timezone.utc) - timedelta(minutes=STALE_AFTER_MINUTES)
    async with AsyncSessionLocal() as db:
        async with db.begin():
            result = await db.execute(
                update(Task)
                .where(
                    Task.status == "running",
                    # Use heartbeat_at when available (more accurate for long tasks),
                    # fall back to started_at for tasks that never wrote a heartbeat.
                    sa.or_(
                        sa.and_(Task.heartbeat_at.isnot(None), Task.heartbeat_at < cutoff),
                        sa.and_(Task.heartbeat_at.is_(None), Task.started_at < cutoff),
                    ),
                )
                .values(status="stale")
                .returning(Task.id)
            )
            stale_ids = result.scalars().all()
    if stale_ids:
        print(f"[startup] Marked {len(stale_ids)} stale task(s): {stale_ids}")
    return len(stale_ids)