from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable

from yuxi.channel.runtime.backoff import BackoffConfig

logger = logging.getLogger(__name__)

StartupTask = Callable[[], Awaitable[Any]]


@dataclass
class StartupTaskResult:
    name: str
    success: bool
    duration_ms: float
    error: str | None = None
    skipped: bool = False


_StartupBackoffConfig = BackoffConfig(
    base_delay=0.5,
    max_delay=5.0,
    exponent=1.5,
    jitter=True,
    jitter_factor=0.1,
    max_retries=3,
)


class StartupTaskRunner:
    """启动任务执行器 — 支持超时、重试、并行执行与结果追踪。"""

    def __init__(self, default_timeout: float = 30.0, max_retries: int = 2, backoff_config: BackoffConfig | None = None):
        self.default_timeout = default_timeout
        self.max_retries = max_retries
        self._backoff_config = backoff_config or _StartupBackoffConfig
        self._results: list[StartupTaskResult] = []
        self._results_lock = asyncio.Lock()

    async def run(
        self,
        task: StartupTask,
        *,
        name: str | None = None,
        timeout: float | None = None,
        retries: int | None = None,
    ) -> StartupTaskResult:
        task_name = name or getattr(task, "__name__", "unknown")
        timeout_val = timeout if timeout is not None else self.default_timeout
        retry_count = retries if retries is not None else self.max_retries

        start = time.time()
        last_error: Exception | None = None

        for attempt in range(retry_count + 1):
            try:
                await asyncio.wait_for(task(), timeout=timeout_val)
                duration_ms = (time.time() - start) * 1000
                result = StartupTaskResult(
                    name=task_name,
                    success=True,
                    duration_ms=duration_ms,
                )
                async with self._results_lock:
                    self._results.append(result)
                logger.info("Startup task '%s' completed in %.2fms", task_name, duration_ms)
                return result
            except asyncio.TimeoutError:
                last_error = asyncio.TimeoutError(f"Task '{task_name}' timed out after {timeout_val}s")
                logger.warning("Startup task '%s' timed out (attempt %d/%d)", task_name, attempt + 1, retry_count + 1)
            except Exception as e:
                last_error = e
                logger.exception("Startup task '%s' failed (attempt %d/%d)", task_name, attempt + 1, retry_count + 1)
                if attempt < retry_count:
                    await asyncio.sleep(self._backoff_config.compute_delay(attempt))

        duration_ms = (time.time() - start) * 1000
        result = StartupTaskResult(
            name=task_name,
            success=False,
            duration_ms=duration_ms,
            error=str(last_error) if last_error else "Unknown error",
        )
        async with self._results_lock:
            self._results.append(result)
        return result

    async def run_parallel(
        self,
        tasks: list[tuple[StartupTask, str]],
        *,
        timeout: float | None = None,
    ) -> list[StartupTaskResult]:
        """并行执行多个启动任务。"""
        coros = [self.run(task, name=name, timeout=timeout) for task, name in tasks]
        return await asyncio.gather(*coros, return_exceptions=True)

    async def get_results(self) -> list[StartupTaskResult]:
        async with self._results_lock:
            return list(self._results)

    async def clear_results(self) -> None:
        async with self._results_lock:
            self._results.clear()