"""Background worker: owns MCP sessions and runs the troubleshooting agent per task.

Consumes task ids from the Redis queue, executes the LangGraph agent, and keeps the
MCP status snapshot fresh for the API/UI.
"""
from __future__ import annotations

import asyncio
import logging

from sqlalchemy import select

from app.agent.graph import run_task_agent
from app.config import settings
from app.database import SessionLocal, init_db
from app.mcp_manager import get_manager
from app.models.enums import TaskStatus
from app.models.task import Task
from app.services.events import TASK_QUEUE, get_redis
from app.services.task_runtime import is_task_cancel_requested

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("worker")


async def _refresh_status_loop() -> None:
    manager = get_manager()
    while True:
        await asyncio.sleep(5)
        try:
            while await manager.process_command_queue_once():
                await asyncio.sleep(0)  # yield so task queue blpop can run
        except Exception as exc:  # noqa: BLE001
            logger.warning("MCP command processing failed: %s", exc)
        try:
            await manager.publish_status()
        except Exception as exc:  # noqa: BLE001
            logger.warning("status publish failed: %s", exc)


async def _task_should_run(task_id: int) -> bool:
    async with SessionLocal() as db:
        res = await db.execute(select(Task).where(Task.id == task_id))
        task = res.scalar_one_or_none()
        if not task:
            return False
        if task.status == TaskStatus.cancelled:
            return False
    if await is_task_cancel_requested(task_id):
        return False
    return True


async def _mark_task_running(task_id: int) -> None:
    async with SessionLocal() as db:
        res = await db.execute(select(Task).where(Task.id == task_id))
        task = res.scalar_one_or_none()
        if task and task.status == TaskStatus.queued:
            task.status = TaskStatus.running
            await db.commit()


async def _run_task(task_id: int, sem: asyncio.Semaphore) -> None:
    async with sem:
        if not await _task_should_run(task_id):
            logger.info("Skipping task %s (cancelled or missing)", task_id)
            return
        logger.info("Running task %s", task_id)
        try:
            await run_task_agent(task_id)
        except asyncio.CancelledError:
            logger.info("task %s asyncio task cancelled", task_id)
        except Exception:  # noqa: BLE001
            logger.exception("task %s crashed", task_id)


async def _cancel_watch_loop(in_flight: dict[int, asyncio.Task]) -> None:
    """Cancel in-flight asyncio tasks when the user requests stop via the API."""
    while True:
        await asyncio.sleep(1)
        for task_id, handle in list(in_flight.items()):
            if handle.done():
                continue
            try:
                if await is_task_cancel_requested(task_id):
                    logger.info("Cancelling in-flight asyncio task %s", task_id)
                    handle.cancel()
            except Exception as exc:  # noqa: BLE001
                logger.debug("cancel watch for task %s: %s", task_id, exc)


async def main() -> None:
    await init_db()
    manager = get_manager()
    logger.info("Starting MCP manager (cloning + launching MCP servers)...")
    await manager.startup()
    logger.info("MCP manager ready: %s", [s["name"] for s in manager.status_snapshot()])

    asyncio.create_task(_refresh_status_loop())

    max_concurrent = max(1, settings.worker_max_concurrent_tasks)
    sem = asyncio.Semaphore(max_concurrent)
    in_flight: dict[int, asyncio.Task] = {}
    asyncio.create_task(_cancel_watch_loop(in_flight))

    redis = get_redis()
    logger.info(
        "Worker waiting for tasks on %s (max concurrent: %s)",
        TASK_QUEUE,
        max_concurrent,
    )
    while True:
        item = await redis.blpop(TASK_QUEUE, timeout=5)
        if not item:
            in_flight = {tid: t for tid, t in in_flight.items() if not t.done()}
            continue
        _, task_id_str = item
        try:
            task_id = int(task_id_str)
        except ValueError:
            continue
        if not await _task_should_run(task_id):
            logger.info("Skipped dequeued task %s (already cancelled)", task_id)
            continue
        logger.info("Picked up task %s", task_id)
        await _mark_task_running(task_id)
        handle = asyncio.create_task(_run_task(task_id, sem))
        in_flight[task_id] = handle

        def _done(t: asyncio.Task, tid: int = task_id) -> None:
            in_flight.pop(tid, None)

        handle.add_done_callback(_done)


if __name__ == "__main__":
    asyncio.run(main())