"""Background task service for asynchronous agent execution.

Manages a queue of background tasks that run agents independently,
with support for concurrent execution, status tracking, and WebSocket notifications.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from ieum.server.services.session_service import Session

logger = logging.getLogger(__name__)


@dataclass
class BackgroundTask:
    """A background agent task."""

    id: str
    session_id: str
    user_id: str
    description: str
    status: str = "queued"  # queued, running, completed, failed, cancelled
    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    started_at: datetime | None = None
    completed_at: datetime | None = None
    result_summary: str | None = None
    error: str | None = None
    branch_name: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "id": self.id,
            "session_id": self.session_id,
            "user_id": self.user_id,
            "description": self.description,
            "status": self.status,
            "created_at": self.created_at.isoformat(),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "result_summary": self.result_summary,
            "error": self.error,
            "branch_name": self.branch_name,
        }


class TaskService:
    """Manages background agent tasks with an asyncio-based queue."""

    def __init__(self, max_concurrent: int = 3) -> None:
        self._max_concurrent = max_concurrent
        self._tasks: dict[str, BackgroundTask] = {}
        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._semaphore = asyncio.Semaphore(max_concurrent)
        self._workers: list[asyncio.Task] = []
        self._running = False
        self._notify_callback: Any = None  # WS broadcast callback

    def set_notify_callback(self, callback: Any) -> None:
        """Set callback for WebSocket notifications."""
        self._notify_callback = callback

    async def start(self, worker_count: int = 2) -> None:
        """Start background worker tasks."""
        if self._running:
            return
        self._running = True
        for _ in range(worker_count):
            worker = asyncio.create_task(self._worker_loop())
            self._workers.append(worker)
        logger.info("TaskService started with %d workers", worker_count)

    async def stop(self) -> None:
        """Stop all workers gracefully."""
        self._running = False
        for worker in self._workers:
            worker.cancel()
        await asyncio.gather(*self._workers, return_exceptions=True)
        self._workers.clear()

    async def create_task(
        self,
        session_id: str,
        user_id: str,
        description: str,
    ) -> BackgroundTask:
        """Create and queue a new background task."""
        task = BackgroundTask(
            id=uuid.uuid4().hex[:12],
            session_id=session_id,
            user_id=user_id,
            description=description,
        )
        self._tasks[task.id] = task
        await self._queue.put(task.id)

        # Notify via WebSocket
        await self._notify(session_id, "task.created", task.to_dict())
        logger.info("Task %s created for session %s", task.id, session_id)
        return task

    def list_tasks(
        self, session_id: str, user_id: str | None = None,
    ) -> list[BackgroundTask]:
        """List tasks for a session, optionally filtered by user."""
        results = []
        for task in self._tasks.values():
            if task.session_id != session_id:
                continue
            if user_id and task.user_id != user_id:
                continue
            results.append(task)
        return sorted(results, key=lambda t: t.created_at, reverse=True)

    def get_task(self, task_id: str) -> BackgroundTask | None:
        """Get a task by ID."""
        return self._tasks.get(task_id)

    async def cancel_task(self, task_id: str) -> bool:
        """Cancel a queued or running task."""
        task = self._tasks.get(task_id)
        if task is None:
            return False
        if task.status in ("completed", "failed", "cancelled"):
            return False
        task.status = "cancelled"
        task.completed_at = datetime.now(UTC)
        await self._notify(task.session_id, "task.cancelled", task.to_dict())
        return True

    async def _worker_loop(self) -> None:
        """Worker loop that processes tasks from the queue."""
        while self._running:
            try:
                task_id = await asyncio.wait_for(self._queue.get(), timeout=1.0)
            except (asyncio.TimeoutError, TimeoutError):
                continue
            except asyncio.CancelledError:
                break

            task = self._tasks.get(task_id)
            if task is None or task.status == "cancelled":
                continue

            async with self._semaphore:
                await self._execute_task(task)

    async def _execute_task(self, task: BackgroundTask) -> None:
        """Execute a single background task."""
        task.status = "running"
        task.started_at = datetime.now(UTC)
        await self._notify(task.session_id, "task.progress", {
            **task.to_dict(),
            "progress": "Starting agent execution...",
        })

        try:
            # Import here to avoid circular imports
            from ieum.server.services.session_service import SessionService

            # Get the session service from the task context
            # The actual agent execution happens through the session's agent
            session = self._get_session(task.session_id)
            if session is None:
                task.status = "failed"
                task.error = "Session not found"
                task.completed_at = datetime.now(UTC)
                await self._notify(task.session_id, "task.failed", task.to_dict())
                return

            if session.agent is None:
                task.status = "failed"
                task.error = "Agent not initialized for session"
                task.completed_at = datetime.now(UTC)
                await self._notify(task.session_id, "task.failed", task.to_dict())
                return

            # Build the execution config
            thread_id = f"bg-{task.id}"  # Separate thread for background task
            config = {
                "configurable": {"thread_id": thread_id},
                "metadata": {
                    "agent_name": session.agent_id,
                    "background_task_id": task.id,
                },
            }

            # Build the prompt
            prompt = (
                f"## Background Task\n\n"
                f"{task.description}\n\n"
                f"Execute this task thoroughly. When done, provide a summary of what was accomplished."
            )

            from langchain_core.messages import HumanMessage

            # Execute the agent
            result = await session.agent.ainvoke(
                {"messages": [HumanMessage(content=prompt)]},
                config=config,
            )

            # Extract summary from the last AI message
            messages = result.get("messages", [])
            summary = ""
            for msg in reversed(messages):
                if hasattr(msg, "content") and isinstance(msg.content, str):
                    summary = msg.content[:500]
                    break

            task.status = "completed"
            task.completed_at = datetime.now(UTC)
            task.result_summary = summary or "Task completed."
            await self._notify(task.session_id, "task.completed", task.to_dict())
            logger.info("Task %s completed", task.id)

        except asyncio.CancelledError:
            task.status = "cancelled"
            task.completed_at = datetime.now(UTC)
            await self._notify(task.session_id, "task.cancelled", task.to_dict())
        except Exception as e:
            task.status = "failed"
            task.error = str(e)[:500]
            task.completed_at = datetime.now(UTC)
            await self._notify(task.session_id, "task.failed", task.to_dict())
            logger.exception("Task %s failed", task.id)

    def _get_session(self, session_id: str) -> Any:
        """Get session from the session service.

        This is set by the app during initialization.
        """
        if hasattr(self, "_session_service") and self._session_service:
            return self._session_service.get_session(session_id)
        return None

    async def _notify(self, session_id: str, event_type: str, data: dict) -> None:
        """Send WebSocket notification."""
        if self._notify_callback:
            try:
                await self._notify_callback(session_id, event_type, data)
            except Exception:
                pass  # Notification is best-effort
