"""Task Scheduler for running collection jobs."""

import asyncio
from datetime import UTC, datetime, timedelta
from typing import Any, Dict, Optional

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger
from sqlalchemy import select

from app.core.logging import get_logger
from app.db.session import async_session_factory
from app.core.time import to_iso8601_utc
from app.models.datasource import DataSource
from app.models.task import CollectionTask
from app.services.collectors.registry import collector_registry

logger = get_logger(__name__)

scheduler = AsyncIOScheduler()
RUNNING_TASK_GUARD_TIMEOUT_MINUTES = 90
RUNNING_COLLECTOR_TASKS: dict[str, asyncio.Task[Any]] = {}


def _collector_task_name(collector_name: str) -> str:
    return f"collector:{collector_name}"


def get_running_collector_task(collector_name: str) -> asyncio.Task[Any] | None:
    task = RUNNING_COLLECTOR_TASKS.get(collector_name)
    if task is not None and not task.done():
        return task

    if task is not None and task.done():
        RUNNING_COLLECTOR_TASKS.pop(collector_name, None)

    target_name = _collector_task_name(collector_name)
    for candidate in asyncio.all_tasks():
        if candidate.done():
            continue
        if candidate.get_name() == target_name:
            RUNNING_COLLECTOR_TASKS[collector_name] = candidate
            return candidate

    return None


async def _update_next_run_at(datasource: DataSource, session) -> None:
    job = scheduler.get_job(datasource.source)
    datasource.next_run_at = job.next_run_time if job else None
    await session.commit()


async def _apply_datasource_schedule(datasource: DataSource, session) -> None:
    collector = collector_registry.get(datasource.source)
    if not collector:
        logger.warning_event(
            "Collector not found for datasource",
            event="collector.schedule.collector_missing",
            context={"collector_name": datasource.source},
        )
        return

    collector_registry.set_active(datasource.source, datasource.is_active)

    existing_job = scheduler.get_job(datasource.source)
    if existing_job:
        scheduler.remove_job(datasource.source)

    if datasource.is_active:
        scheduler.add_job(
            run_collector_task,
            trigger=IntervalTrigger(minutes=max(1, datasource.frequency_minutes)),
            id=datasource.source,
            name=datasource.name,
            replace_existing=True,
            kwargs={"collector_name": datasource.source},
        )
        logger.info_event(
            "Scheduled collector",
            event="collector.schedule.updated",
            context={"collector_name": datasource.source, "frequency_minutes": datasource.frequency_minutes},
        )
    else:
        logger.info_event(
            "Collector disabled",
            event="collector.schedule.disabled",
            context={"collector_name": datasource.source},
        )

    await _update_next_run_at(datasource, session)


async def run_collector_task(collector_name: str):
    """Run a single collector task."""
    collector = collector_registry.get(collector_name)
    if not collector:
        logger.error_event(
            "Collector not found",
            event="collector.run.collector_missing",
            context={"collector_name": collector_name},
        )
        return

    async with async_session_factory() as db:
        result = await db.execute(select(DataSource).where(DataSource.source == collector_name))
        datasource = result.scalar_one_or_none()
        if not datasource:
            logger.error_event(
                "Datasource not found for collector",
                event="collector.run.datasource_missing",
                context={"collector_name": collector_name},
            )
            return

        if not datasource.is_active:
            logger.info_event(
                "Skipping disabled collector",
                event="collector.run.skipped_disabled",
                context={"collector_name": collector_name},
            )
            return

        running_result = await db.execute(
            select(CollectionTask)
            .where(
                CollectionTask.datasource_id == datasource.id,
                CollectionTask.status == "running",
            )
            .order_by(CollectionTask.started_at.desc(), CollectionTask.id.desc())
            .limit(1)
        )
        existing_running = running_result.scalar_one_or_none()
        if existing_running is not None:
            now = datetime.now(UTC)
            started_at = existing_running.started_at
            if started_at is not None and started_at.tzinfo is None:
                started_at = started_at.replace(tzinfo=UTC)

            is_stale = (
                started_at is not None
                and (now - started_at) > timedelta(minutes=RUNNING_TASK_GUARD_TIMEOUT_MINUTES)
            )
            if not is_stale:
                logger.warning_event(
                    "Skipping collector trigger because task is already running",
                    event="collector.run.skipped_already_running",
                    context={"collector_name": collector_name, "task_id": existing_running.id},
                )
                return

            existing_error = (existing_running.error_message or "").strip()
            stale_reason = (
                f"Marked failed automatically after stale running timeout "
                f"({RUNNING_TASK_GUARD_TIMEOUT_MINUTES}m) in scheduler guard"
            )
            existing_running.status = "failed"
            existing_running.phase = "failed"
            existing_running.completed_at = now
            existing_running.error_message = (
                f"{existing_error}\n{stale_reason}".strip()
                if existing_error
                else stale_reason
            )
            await db.commit()
            logger.warning_event(
                "Marked stale running task as failed before rerun",
                event="collector.run.stale_task_failed",
                context={"collector_name": collector_name, "task_id": existing_running.id},
            )

        try:
            collector._datasource_id = datasource.id
            logger.info_event(
                "Running collector",
                event="collector.run.started",
                context={"collector_name": collector_name, "datasource_id": datasource.id},
            )
            task_result = await collector.run(db)
            datasource.last_run_at = datetime.now(UTC)
            datasource.last_status = task_result.get("status")
            await _update_next_run_at(datasource, db)
            logger.info_event(
                "Collector completed",
                event="collector.run.completed",
                context={"collector_name": collector_name, "datasource_id": datasource.id, "result": task_result},
            )
        except asyncio.CancelledError:
            datasource.last_run_at = datetime.now(UTC)
            datasource.last_status = "cancelled"
            await db.commit()
            logger.warning_event(
                "Collector cancelled by operator",
                event="collector.run.cancelled",
                context={"collector_name": collector_name, "datasource_id": datasource.id},
            )
            raise
        except Exception as exc:
            datasource.last_run_at = datetime.now(UTC)
            datasource.last_status = "failed"
            await db.commit()
            logger.exception_event(
                "Collector failed",
                event="collector.run.failed",
                context={"collector_name": collector_name, "datasource_id": datasource.id, "error": str(exc)},
            )


async def cleanup_stale_running_tasks(max_age_hours: int = 2) -> int:
    """Mark stale running tasks as failed after restarts or collector hangs."""
    cutoff = datetime.now(UTC) - timedelta(hours=max_age_hours)

    async with async_session_factory() as db:
        result = await db.execute(
            select(CollectionTask).where(
                CollectionTask.status == "running",
                CollectionTask.started_at.is_not(None),
                CollectionTask.started_at < cutoff,
            )
        )
        stale_tasks = result.scalars().all()

        for task in stale_tasks:
            task.status = "failed"
            task.phase = "failed"
            task.completed_at = datetime.now(UTC)
            existing_error = (task.error_message or "").strip()
            cleanup_error = "Marked failed automatically after stale running task cleanup"
            task.error_message = f"{existing_error}\n{cleanup_error}".strip() if existing_error else cleanup_error

        if stale_tasks:
            await db.commit()
            logger.warning_event(
                "Cleaned up stale running collection tasks",
                event="collector.cleanup.stale_tasks_cleaned",
                context={"count": len(stale_tasks)},
            )

        return len(stale_tasks)


def start_scheduler() -> None:
    """Start the scheduler."""
    if not scheduler.running:
        scheduler.start()
        logger.info_event("Scheduler started", event="scheduler.started")


def stop_scheduler() -> None:
    """Stop the scheduler."""
    if scheduler.running:
        scheduler.shutdown(wait=False)
        logger.info_event("Scheduler stopped", event="scheduler.stopped")


async def sync_scheduler_with_datasources() -> None:
    """Synchronize scheduler jobs with datasource table."""
    async with async_session_factory() as db:
        result = await db.execute(select(DataSource).order_by(DataSource.id))
        datasources = result.scalars().all()

        configured_sources = {datasource.source for datasource in datasources}
        for job in list(scheduler.get_jobs()):
            if job.id not in configured_sources:
                scheduler.remove_job(job.id)

        for datasource in datasources:
            await _apply_datasource_schedule(datasource, db)


async def sync_datasource_job(datasource_id: int) -> bool:
    """Synchronize a single datasource job after settings changes."""
    async with async_session_factory() as db:
        datasource = await db.get(DataSource, datasource_id)
        if not datasource:
            return False

        await _apply_datasource_schedule(datasource, db)
        return True


def get_scheduler_jobs() -> list[Dict[str, Any]]:
    """Get all scheduled jobs."""
    jobs = []
    for job in scheduler.get_jobs():
        jobs.append(
            {
                "id": job.id,
                "name": job.name,
                "next_run_time": to_iso8601_utc(job.next_run_time),
                "trigger": str(job.trigger),
            }
        )
    return jobs


async def get_latest_task_id_for_datasource(datasource_id: int) -> Optional[int]:
    from app.models.task import CollectionTask

    async with async_session_factory() as db:
        result = await db.execute(
            select(CollectionTask.id)
            .where(CollectionTask.datasource_id == datasource_id)
            .order_by(CollectionTask.created_at.desc(), CollectionTask.id.desc())
            .limit(1)
        )
        return result.scalar_one_or_none()


def run_collector_now(collector_name: str) -> bool:
    """Run a collector immediately (not scheduled)."""
    collector = collector_registry.get(collector_name)
    if not collector:
        logger.error_event(
            "Collector not found",
            event="collector.trigger.collector_missing",
            context={"collector_name": collector_name},
        )
        return False

    existing_task = get_running_collector_task(collector_name)
    if existing_task is not None and not existing_task.done():
        logger.warning_event(
            "Collector is already running in-memory; skipping duplicate trigger",
            event="collector.trigger.skipped_already_running",
            context={"collector_name": collector_name},
        )
        return False

    try:
        task = asyncio.create_task(run_collector_task(collector_name), name=_collector_task_name(collector_name))
        RUNNING_COLLECTOR_TASKS[collector_name] = task

        def _cleanup_task(done_task: asyncio.Task[Any]) -> None:
            current = RUNNING_COLLECTOR_TASKS.get(collector_name)
            if current is done_task:
                RUNNING_COLLECTOR_TASKS.pop(collector_name, None)

        task.add_done_callback(_cleanup_task)
        logger.info_event(
            "Triggered collector",
            event="collector.trigger.started",
            context={"collector_name": collector_name},
        )
        return True
    except Exception as exc:
        logger.error_event(
            "Failed to trigger collector",
            event="collector.trigger.failed",
            context={"collector_name": collector_name, "error": str(exc)},
        )
        return False


async def cancel_running_collector_now(collector_name: str) -> bool:
    task = get_running_collector_task(collector_name)
    if task is None or task.done():
        RUNNING_COLLECTOR_TASKS.pop(collector_name, None)
        return False

    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        return True
    return task.cancelled()