import logging
import uuid
from collections.abc import Awaitable, Callable, Sequence
from datetime import UTC, datetime

from ai.backend.common.exception import BackendAIError, ErrorCode
from ai.backend.logging.utils import BraceStyleAdapter
from ai.backend.manager.actions.action import (
    BaseAction,
    BaseActionResult,
    BaseActionResultMeta,
    BaseActionTriggerMeta,
    ProcessResult,
)
from ai.backend.manager.actions.monitors.monitor import ActionMonitor
from ai.backend.manager.actions.types import OperationStatus
from ai.backend.manager.actions.validator.base import ActionValidator

log = BraceStyleAdapter(logging.getLogger(__spec__.name))


class ActionRunner[TAction: BaseAction, TActionResult: BaseActionResult]:
    _func: Callable[[TAction], Awaitable[TActionResult]]
    _monitors: Sequence[ActionMonitor]

    def __init__(
        self,
        func: Callable[[TAction], Awaitable[TActionResult]],
        monitors: Sequence[ActionMonitor] | None,
    ) -> None:
        self._func = func
        self._monitors = monitors or []

    async def _start_monitors(
        self, action: TAction, action_trigger_meta: BaseActionTriggerMeta
    ) -> None:
        for monitor in self._monitors:
            try:
                await monitor.prepare(action, action_trigger_meta)
            except Exception as e:
                log.warning("Error in monitor prepare method: {}", e)

    async def _finalize_monitors(
        self,
        action: TAction,
        meta: BaseActionResultMeta,
    ) -> None:
        process_result = ProcessResult(meta=meta)
        for monitor in reversed(self._monitors):
            try:
                await monitor.done(action, process_result)
            except Exception as e:
                log.warning("Error in monitor done method: {}", e)

    async def run(
        self, action: TAction, action_trigger_meta: BaseActionTriggerMeta
    ) -> TActionResult:
        started_at = action_trigger_meta.started_at
        action_id = action_trigger_meta.action_id
        status = OperationStatus.UNKNOWN
        description: str = "unknown"
        result: TActionResult | None = None
        error_code: ErrorCode | None = None

        await self._start_monitors(action, action_trigger_meta)
        try:
            result = await self._func(action)
        except BackendAIError as e:
            log.exception("Action processing error: {}", e)
            status = OperationStatus.ERROR
            description = str(e)
            error_code = e.error_code()
            raise
        except BaseException as e:
            log.exception("Unexpected error during action processing: {}", e)
            status = OperationStatus.ERROR
            description = str(e)
            error_code = ErrorCode.default()
            raise
        else:
            status = OperationStatus.SUCCESS
            description = "Success"
            return result
        finally:
            entity_id = action.entity_id()
            if entity_id is None and result is not None:
                entity_id = result.entity_id()
            ended_at = datetime.now(UTC)
            duration = ended_at - started_at
            meta = BaseActionResultMeta(
                action_id=action_id,
                entity_id=entity_id,
                status=status,
                description=description,
                started_at=started_at,
                ended_at=ended_at,
                duration=duration,
                error_code=error_code,
            )
            await self._finalize_monitors(
                action,
                meta,
            )


class ActionProcessor[TAction: BaseAction, TActionResult: BaseActionResult]:
    _validators: Sequence[ActionValidator]

    _runner: ActionRunner[TAction, TActionResult]

    def __init__(
        self,
        func: Callable[[TAction], Awaitable[TActionResult]],
        monitors: Sequence[ActionMonitor] | None = None,
        validators: Sequence[ActionValidator] | None = None,
    ) -> None:
        self._runner = ActionRunner(func, monitors)

        self._validators = validators or []

    async def _run(self, action: TAction) -> TActionResult:
        started_at = datetime.now(UTC)
        action_id = uuid.uuid4()
        action_trigger_meta = BaseActionTriggerMeta(action_id=action_id, started_at=started_at)
        for validator in self._validators:
            await validator.validate(action, action_trigger_meta)

        return await self._runner.run(action, action_trigger_meta)

    async def wait_for_complete(self, action: TAction) -> TActionResult:
        return await self._run(action)
