"""
TaskService — creates tasks, dispatches them as asyncio background tasks,
and persists every state change to the DB so restarts don't lose history.

Flow:
  create() → INSERT tasks (status=pending) → asyncio.create_task(run) → return task_id
  run()    → UPDATE running → per-device execute → UPDATE success/failed
"""

import asyncio
from datetime import datetime, timezone
from uuid import uuid4

from sqlalchemy import select
from sqlalchemy.orm import selectinload

from app.models.device import Device
from app.models.object import Object
from app.models.task import Task, TaskEvent, TaskResult
from app.models.user import User
from app.plugins.registry import ActionRegistry, action_registry
from app.services.database import AsyncSessionLocal
from app.services.sse import SSEEvent, sse_manager

# Max parallel devices within a single task
_TASK_DEVICE_CONCURRENCY = 20


async def _emit(task_id: str, message: str, level: str = "info") -> None:
    """Write a progress line to DB and push to SSE subscribers."""
    now = datetime.now(timezone.utc)
    async with AsyncSessionLocal() as db:
        async with db.begin():
            db.add(TaskEvent(task_id=task_id, message=message, level=level, created_at=now))

    await sse_manager.publish(SSEEvent(
        type="task_event",
        task_id=task_id,
        data={"message": message, "level": level, "ts": now.isoformat()},
    ))


async def _resolve_scope(scope: dict) -> list[tuple[Device, Object]]:
    """Expand target_scope to a list of (device, parent_object) pairs."""
    async with AsyncSessionLocal() as db:
        query = select(Device).options(selectinload(Device.object)).where(Device.is_active == True)

        scope_type = scope.get("type")
        if scope_type == "device":
            query = query.where(Device.id == scope["id"])
        elif scope_type == "object":
            query = query.where(Device.object_id == scope["id"])
        elif scope_type == "category":
            query = query.where(Device.category == scope["value"])
        elif scope_type == "role":
            query = query.where(Device.role == scope["value"])
            if "object_id" in scope:
                query = query.where(Device.object_id == scope["object_id"])
        else:
            return []

        result = await db.execute(query)
        devices = result.scalars().all()
        return [(d, d.object) for d in devices]


async def _run_on_device(
    task_id: str,
    device: Device,
    obj: Object,
    action,
    params: dict,
    sem: asyncio.Semaphore,
) -> bool:
    async with sem:
        emit = lambda msg, level="info": _emit(task_id, msg, level)

        # Create task_result row
        result_id: int | None = None
        started_at = datetime.now(timezone.utc)
        async with AsyncSessionLocal() as db:
            async with db.begin():
                tr = TaskResult(
                    task_id=task_id,
                    device_id=device.id,
                    status="running",
                    executed_at=started_at,
                )
                db.add(tr)
                await db.flush()
                result_id = tr.id

        try:
            action_result = await action.execute(device, obj, params, emit)
        except Exception as exc:
            action_result_success = False
            action_stdout = ""
            action_stderr = ""
            action_exit_code = -1
            action_data = {}
            action_error = str(exc)
        else:
            action_result_success = action_result.success
            action_stdout = action_result.stdout
            action_stderr = action_result.stderr
            action_exit_code = action_result.exit_code
            action_data = action_result.data
            action_error = action_result.error

        finished_at = datetime.now(timezone.utc)
        duration_ms = int((finished_at - started_at).total_seconds() * 1000)

        async with AsyncSessionLocal() as db:
            async with db.begin():
                tr = await db.get(TaskResult, result_id)
                if tr:
                    tr.status = "success" if action_result_success else "failed"
                    tr.stdout = action_stdout
                    tr.stderr = action_stderr
                    tr.exit_code = action_exit_code
                    tr.data = action_data or None
                    tr.duration_ms = duration_ms

        await sse_manager.publish(SSEEvent(
            type="task_result",
            task_id=task_id,
            data={
                "device_id": device.id,
                "device_ip": device.ip,
                "status": "success" if action_result_success else "failed",
                "duration_ms": duration_ms,
                "error": action_error or None,
            },
        ))

        return action_result_success


async def _run_task(task_id: str) -> None:
    await _emit(task_id, "Task started", "info")

    async with AsyncSessionLocal() as db:
        async with db.begin():
            task = await db.get(Task, task_id)
            if task is None:
                return
            task.status = "running"
            task.started_at = datetime.now(timezone.utc)
            scope = dict(task.target_scope)
            action_type = task.type
            params = dict(task.params)

    action = action_registry.get(action_type)
    if action is None:
        async with AsyncSessionLocal() as db:
            async with db.begin():
                task = await db.get(Task, task_id)
                if task:
                    task.status = "failed"
                    task.error_message = f"Unknown action: {action_type}"
                    task.finished_at = datetime.now(timezone.utc)
        await _emit(task_id, f"Unknown action: {action_type}", "error")
        return

    targets = await _resolve_scope(scope)
    if not targets:
        async with AsyncSessionLocal() as db:
            async with db.begin():
                task = await db.get(Task, task_id)
                if task:
                    task.status = "success"
                    task.result = {"total": 0, "ok": 0, "failed": 0}
                    task.finished_at = datetime.now(timezone.utc)
        await _emit(task_id, "No devices matched the scope", "warn")
        return

    await _emit(task_id, f"Running on {len(targets)} device(s)...", "info")
    sem = asyncio.Semaphore(_TASK_DEVICE_CONCURRENCY)

    results = await asyncio.gather(
        *[_run_on_device(task_id, device, obj, action, params, sem) for device, obj in targets],
        return_exceptions=True,
    )

    ok = sum(1 for r in results if r is True)
    failed = len(results) - ok
    final_status = "success" if failed == 0 else ("failed" if ok == 0 else "partial")

    async with AsyncSessionLocal() as db:
        async with db.begin():
            task = await db.get(Task, task_id)
            if task:
                task.status = final_status
                task.result = {"total": len(results), "ok": ok, "failed": failed}
                task.finished_at = datetime.now(timezone.utc)

    await _emit(task_id, f"Done: {ok}/{len(results)} succeeded", "info")
    await sse_manager.publish(SSEEvent(
        type="task_status",
        task_id=task_id,
        data={"status": final_status, "total": len(results), "ok": ok, "failed": failed},
    ))


async def create_task(
    action_type: str,
    target_scope: dict,
    params: dict,
    user_id: int | None = None,
) -> Task:
    if action_registry.get(action_type) is None:
        raise ValueError(f"Unknown action: {action_type}")

    now = datetime.now(timezone.utc)
    task_id = str(uuid4())

    async with AsyncSessionLocal() as db:
        async with db.begin():
            task = Task(
                id=task_id,
                type=action_type,
                status="pending",
                target_scope=target_scope,
                params=params,
                created_by=user_id,
                created_at=now,
            )
            db.add(task)

    # Dispatch — fire and forget; state is already persisted
    asyncio.create_task(_run_task(task_id))

    # Return a fresh copy to the caller
    async with AsyncSessionLocal() as db:
        result = await db.get(Task, task_id)
        return result