from __future__ import annotations

import asyncio
import time
import uuid
from collections.abc import Coroutine
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from relace_toolbox.types import ToolResult


class TaskStatus(str, Enum):
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMEOUT = "timeout"


@dataclass
class TaskResult:
    status: TaskStatus
    result: ToolResult | None
    error: str | None
    elapsed_seconds: float


@dataclass
class AsyncTask:
    task: asyncio.Task[ToolResult]
    started_at: float


class AsyncTaskManager:
    """Manages async tool execution tasks.

    Tools can submit long-running operations that may be suspended if they
    exceed a wait_before_suspend timeout. The manager tracks these tasks and
    allows resuming the wait via wait_for_task.

    All tasks have a system-controlled cancel_after timeout (total wall time)
    after which they are cancelled.
    """

    def __init__(self) -> None:
        self._tasks: dict[str, AsyncTask] = {}

    async def submit_task(
        self,
        coro: Coroutine[Any, Any, ToolResult],
        cancel_after: float,
    ) -> str:
        """Submit a coroutine for async execution."""
        task_id = str(uuid.uuid4())
        task = asyncio.create_task(asyncio.wait_for(coro, timeout=cancel_after))
        self._tasks[task_id] = AsyncTask(
            task=task,
            started_at=time.time(),
        )
        return task_id

    async def wait_for_task(
        self,
        task_id: str,
        wait_before_suspend: float,
    ) -> TaskResult:
        """Wait for a task to complete.

        If task doesn't complete within wait_before_suspend, returns RUNNING status.
        """
        t = self._tasks.get(task_id)
        if t is None:
            raise ValueError(f"Unknown task: {task_id}")

        # Wait for task to complete or timeout
        if not t.task.done():
            await asyncio.wait({t.task}, timeout=wait_before_suspend)

        elapsed = time.time() - t.started_at

        if not t.task.done():
            return TaskResult(TaskStatus.RUNNING, None, None, elapsed)

        # Task completed - remove from tracking and return result
        del self._tasks[task_id]

        if t.task.cancelled():
            return TaskResult(TaskStatus.TIMEOUT, None, "Task was cancelled", elapsed)

        exc = t.task.exception()
        if exc is not None:
            if isinstance(exc, TimeoutError):
                return TaskResult(
                    TaskStatus.TIMEOUT,
                    None,
                    f"Task timed out after {elapsed:.1f}s",
                    elapsed,
                )
            return TaskResult(TaskStatus.FAILED, None, str(exc), elapsed)

        return TaskResult(TaskStatus.COMPLETED, t.task.result(), None, elapsed)
