"""Agent execution service."""

import asyncio
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
    from ieum.server.websocket.adapter import WebSocketUIAdapter


@dataclass
class ExecutionContext:
    """Context for agent execution."""

    session: Any  # Session
    adapter: Any  # WebSocketUIAdapter
    task: asyncio.Task | None = None
    started_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    cancelled: bool = False


class AgentService:
    """Handles agent execution for WebSocket sessions."""

    def __init__(self) -> None:
        """Initialize the agent service."""
        self._executions: dict[str, ExecutionContext] = {}
        self._lock = asyncio.Lock()

    async def execute(
        self,
        session: Any,  # Session
        user_input: str,
        adapter: Any,  # WebSocketUIAdapter
    ) -> None:
        """Execute a task with the agent.

        Args:
            session: The session to execute in.
            user_input: User's input message.
            adapter: WebSocket UI adapter for streaming output.
        """
        from ieum.server.services.session_service import SessionService
        from ieum.server.websocket.adapter import execute_task_websocket

        # Initialize agent if needed
        service = SessionService()
        await service.initialize_agent(session)

        if session.agent is None:
            await adapter.send_error("Failed to initialize agent")
            return

        # Create execution context
        ctx = ExecutionContext(session=session, adapter=adapter)

        async with self._lock:
            # Cancel existing execution for this session
            existing = self._executions.get(session.id)
            if existing and existing.task and not existing.task.done():
                existing.task.cancel()
                existing.cancelled = True

            self._executions[session.id] = ctx

        # Create and store the task
        task = asyncio.create_task(
            execute_task_websocket(
                user_input=user_input,
                agent=session.agent,
                assistant_id=session.agent_id,
                session=session,
                adapter=adapter,
                backend=session.backend,
            )
        )
        ctx.task = task

        try:
            await task
        except asyncio.CancelledError:
            await adapter.send_system_message("Task cancelled")
        except Exception as e:
            await adapter.send_error(str(e))
        finally:
            async with self._lock:
                self._executions.pop(session.id, None)

    async def cancel(self, session_id: str) -> bool:
        """Cancel an ongoing execution.

        Args:
            session_id: Session ID to cancel.

        Returns:
            True if cancelled, False if no execution found.
        """
        async with self._lock:
            ctx = self._executions.get(session_id)
            if ctx and ctx.task and not ctx.task.done():
                ctx.task.cancel()
                ctx.cancelled = True
                return True
            return False

    def is_executing(self, session_id: str) -> bool:
        """Check if a session is currently executing.

        Args:
            session_id: Session ID to check.

        Returns:
            True if executing, False otherwise.
        """
        ctx = self._executions.get(session_id)
        return ctx is not None and ctx.task is not None and not ctx.task.done()
