from __future__ import annotations

import asyncio
import inspect
import logging
from collections import Counter, defaultdict, deque
from dataclasses import dataclass
from time import monotonic
from typing import Any, Generic, TypeVar
from uuid import uuid4

from pydantic import TypeAdapter, ValidationError

from .context import WorkflowContext
from .errors import (
    PayloadValidationError,
    StateValidationError,
    TaskTimeoutError,
    WorkflowExecutionError,
    WorkflowLimitExceededError,
    WorkflowTimeoutError,
)
from .history import HistoryRecorder
from .models import Event, RunResult, RunStatus, StateT, TaskSpec, INTERNAL_START_EVENT

if False:  # pragma: no cover
    from .definition import Workflow


@dataclass
class _RunCounters:
    processed_events: int = 0
    task_invocations: int = 0
    invocations_per_task: Counter[str] = None  # type: ignore[assignment]

    def __post_init__(self):
        if self.invocations_per_task is None:
            self.invocations_per_task = Counter()


@dataclass
class _TaskOutcome(Generic[StateT]):
    state: StateT
    emitted_events: list[Event]
    stop_requested: bool = False
    stop_result: Any = None
    failure_status: RunStatus | None = None
    failure_reason: str | None = None


class WorkflowEngine(Generic[StateT]):
    def __init__(self, workflow: "Workflow[StateT]"):
        self.workflow = workflow
        self.logger = logging.getLogger(
            f"{__name__}.{self.workflow.__class__.__name__}"
        )
        self.task_map = {spec.name: spec for spec in workflow.__workflow_tasks__}
        self.router = defaultdict(list)
        for route in workflow.__workflow_routes__:
            self.router[route.event].append(self.task_map[route.task])

    async def run(
        self,
        payload: Any = None,
        *,
        state: StateT | None = None,
        run_id: str | None = None,
    ) -> RunResult[StateT]:
        run_id = run_id or uuid4().hex
        recorder = HistoryRecorder(run_id=run_id, workflow_name=self.workflow.workflow_name)
        counters = _RunCounters()
        run_started_at = monotonic()

        try:
            committed_state = (
                self.workflow._state_adapter.validate(state)
                if state is not None
                else self.workflow._make_initial_state()
            )
        except StateValidationError as exc:
            recorder.record(
                "workflow_failed",
                details={"reason": str(exc), "error_type": exc.__class__.__name__},
            )
            return RunResult(
                run_id=run_id,
                status=RunStatus.FAILED,
                result=None,
                state=state,  # type: ignore[arg-type]
                history=recorder.records,
                termination_reason=str(exc),
            )

        queue = deque(
            [
                Event(
                    event_id=uuid4().hex,
                    name=INTERNAL_START_EVENT,
                    payload=payload,
                )
            ]
        )

        recorder.record(
            "workflow_started",
            details={"initial_payload": payload},
        )
        recorder.record(
            "event_queued",
            event_id=queue[0].event_id,
            event_name=queue[0].name,
            details={"source_task": None},
        )

        status = RunStatus.COMPLETED
        result = None
        termination_reason = None

        try:
            while queue:
                self._enforce_run_timeout(run_started_at)
                event = queue.popleft()
                counters.processed_events += 1
                self._enforce_event_limit(counters)

                recorder.record(
                    "event_dequeued",
                    event_id=event.event_id,
                    event_name=event.name,
                    details={"source_task": event.source_task},
                )

                subscribers = self.router.get(event.name, ())
                if not subscribers:
                    recorder.record(
                        "event_unhandled",
                        event_id=event.event_id,
                        event_name=event.name,
                    )
                    continue

                for task_spec in subscribers:
                    self._enforce_run_timeout(run_started_at)
                    outcome = await self._invoke_task(
                        task_spec=task_spec,
                        event=event,
                        state=committed_state,
                        run_started_at=run_started_at,
                        counters=counters,
                        recorder=recorder,
                    )
                    committed_state = outcome.state

                    if outcome.failure_status is not None:
                        status = outcome.failure_status
                        termination_reason = outcome.failure_reason
                        queue.clear()
                        break

                    if outcome.stop_requested:
                        status = RunStatus.STOPPED
                        result = outcome.stop_result
                        termination_reason = "stopped by task"
                        queue.clear()
                        break

                    for emitted in outcome.emitted_events:
                        queue.append(emitted)

                if status is not RunStatus.COMPLETED:
                    break

        except WorkflowLimitExceededError as exc:
            status = RunStatus.LIMIT_REACHED
            termination_reason = str(exc)
            recorder.record(
                "workflow_stopped",
                details={"reason": termination_reason, "status": status.value},
            )
        except WorkflowTimeoutError as exc:
            status = RunStatus.TIMEOUT
            termination_reason = str(exc)
            recorder.record(
                "workflow_stopped",
                details={"reason": termination_reason, "status": status.value},
            )
        else:
            record_type = "workflow_completed"
            if status == RunStatus.STOPPED:
                record_type = "workflow_stopped"
            elif status == RunStatus.FAILED:
                record_type = "workflow_failed"
            recorder.record(
                record_type,
                details={"reason": termination_reason, "result": result},
            )

        return RunResult(
            run_id=run_id,
            status=status,
            result=result,
            state=committed_state,
            history=recorder.records,
            termination_reason=termination_reason,
        )

    async def _invoke_task(
        self,
        *,
        task_spec: TaskSpec,
        event: Event,
        state: StateT,
        run_started_at: float,
        counters: _RunCounters,
        recorder: HistoryRecorder,
    ) -> _TaskOutcome[StateT]:
        retry = task_spec.retry

        for attempt in range(1, retry.max_attempts + 1):
            self._enforce_run_timeout(run_started_at)
            counters.task_invocations += 1
            counters.invocations_per_task[task_spec.name] += 1
            self._enforce_task_limits(task_spec.name, counters)

            working_state = self.workflow._state_adapter.clone(state)
            context = WorkflowContext(
                run_id=recorder.run_id,
                workflow_name=recorder.workflow_name,
                task_name=task_spec.name,
                event=event,
                attempt=attempt,
                state=working_state,
                history=recorder,
                logger=self.logger,
            )
            recorder.record(
                "task_started",
                task_name=task_spec.name,
                event_id=event.event_id,
                event_name=event.name,
                attempt=attempt,
            )

            try:
                validated_payload = self._validate_payload(task_spec, event.payload)
                result = await self._call_handler(task_spec, context, validated_payload)
                committed_state = self.workflow._state_adapter.validate(working_state)
            except (PayloadValidationError, StateValidationError) as exc:
                recorder.record(
                    "task_failed",
                    task_name=task_spec.name,
                    event_id=event.event_id,
                    event_name=event.name,
                    attempt=attempt,
                    details={
                        "error_type": exc.__class__.__name__,
                        "error_message": str(exc),
                    },
                )
                return await self._handle_terminal_task_error(
                    task_spec=task_spec,
                    state=state,
                    event=event,
                    attempt=attempt,
                    error=exc,
                    recorder=recorder,
                )
            except TaskTimeoutError as exc:
                recorder.record(
                    "task_failed",
                    task_name=task_spec.name,
                    event_id=event.event_id,
                    event_name=event.name,
                    attempt=attempt,
                    details={
                        "error_type": exc.__class__.__name__,
                        "error_message": str(exc),
                    },
                )
                if retry.should_retry(exc, attempt):
                    await self._schedule_retry(
                        task_spec=task_spec,
                        attempt=attempt,
                        error=exc,
                        recorder=recorder,
                    )
                    continue
                return await self._handle_terminal_task_error(
                    task_spec=task_spec,
                    state=state,
                    event=event,
                    attempt=attempt,
                    error=exc,
                    recorder=recorder,
                )
            except Exception as exc:  # noqa: BLE001
                recorder.record(
                    "task_failed",
                    task_name=task_spec.name,
                    event_id=event.event_id,
                    event_name=event.name,
                    attempt=attempt,
                    details={
                        "error_type": exc.__class__.__name__,
                        "error_message": str(exc),
                    },
                )
                if retry.should_retry(exc, attempt):
                    await self._schedule_retry(
                        task_spec=task_spec,
                        attempt=attempt,
                        error=exc,
                        recorder=recorder,
                    )
                    continue
                return await self._handle_terminal_task_error(
                    task_spec=task_spec,
                    state=state,
                    event=event,
                    attempt=attempt,
                    error=exc,
                    recorder=recorder,
                )

            recorder.record(
                "task_succeeded",
                task_name=task_spec.name,
                event_id=event.event_id,
                event_name=event.name,
                attempt=attempt,
                details={"return_value": result},
            )

            if context.stop_requested:
                return _TaskOutcome(
                    state=committed_state,
                    emitted_events=[],
                    stop_requested=True,
                    stop_result=context.stop_result,
                )

            for emitted in context.outbox:
                recorder.record(
                    "event_queued",
                    task_name=task_spec.name,
                    event_id=emitted.event_id,
                    event_name=emitted.name,
                    attempt=attempt,
                    details={"source_task": emitted.source_task},
                )

            return _TaskOutcome(
                state=committed_state,
                emitted_events=context.outbox,
            )

        return _TaskOutcome(
            state=state,
            emitted_events=[],
            failure_status=RunStatus.FAILED,
            failure_reason=f"Task {task_spec.name} exceeded retry attempts",
        )

    async def _handle_terminal_task_error(
        self,
        *,
        task_spec: TaskSpec,
        state: StateT,
        event: Event,
        attempt: int,
        error: BaseException,
        recorder: HistoryRecorder,
    ) -> _TaskOutcome[StateT]:
        if task_spec.failure_event:
            failure_payload = {
                "task": task_spec.name,
                "event": event.name,
                "attempt": attempt,
                "error_type": error.__class__.__name__,
                "error_message": str(error),
                "original_payload": event.payload,
            }
            emitted = Event(
                event_id=uuid4().hex,
                name=task_spec.failure_event,
                payload=failure_payload,
                source_task=task_spec.name,
                source_event_id=event.event_id,
            )
            recorder.record(
                "event_queued",
                task_name=task_spec.name,
                event_id=emitted.event_id,
                event_name=emitted.name,
                attempt=attempt,
                details={"source_task": emitted.source_task, "failure_event": True},
            )
            return _TaskOutcome(state=state, emitted_events=[emitted])

        return _TaskOutcome(
            state=state,
            emitted_events=[],
            failure_status=RunStatus.FAILED,
            failure_reason=f"{task_spec.name} failed: {error}",
        )

    async def _schedule_retry(
        self,
        *,
        task_spec: TaskSpec,
        attempt: int,
        error: BaseException,
        recorder: HistoryRecorder,
    ) -> None:
        delay = task_spec.retry.delay_for_attempt(attempt)
        recorder.record(
            "retry_scheduled",
            task_name=task_spec.name,
            attempt=attempt,
            details={
                "delay_seconds": delay,
                "error_type": error.__class__.__name__,
                "error_message": str(error),
            },
        )
        if delay > 0:
            await asyncio.sleep(delay)

    def _validate_payload(self, task_spec: TaskSpec, payload: Any) -> Any:
        if task_spec.input_type is None:
            return payload
        try:
            return TypeAdapter(task_spec.input_type).validate_python(payload)
        except ValidationError as exc:
            raise PayloadValidationError(str(exc)) from exc

    async def _call_handler(
        self,
        task_spec: TaskSpec,
        context: WorkflowContext[StateT],
        payload: Any,
    ) -> Any:
        handler = getattr(self.workflow, task_spec.method_name)
        timeout = task_spec.timeout

        try:
            if inspect.iscoroutinefunction(handler):
                coro = handler(context, payload)
                if timeout is None:
                    return await coro
                return await asyncio.wait_for(coro, timeout=timeout)

            if timeout is None:
                return handler(context, payload)

            return await asyncio.wait_for(
                asyncio.to_thread(handler, context, payload),
                timeout=timeout,
            )
        except TimeoutError as exc:
            raise TaskTimeoutError(
                f"Task {task_spec.name} timed out after {timeout} seconds"
            ) from exc

    def _enforce_event_limit(self, counters: _RunCounters) -> None:
        if counters.processed_events > self.workflow.policies.max_events_processed:
            raise WorkflowLimitExceededError("Workflow exceeded max_events_processed")

    def _enforce_task_limits(self, task_name: str, counters: _RunCounters) -> None:
        if counters.task_invocations > self.workflow.policies.max_task_invocations:
            raise WorkflowLimitExceededError("Workflow exceeded max_task_invocations")
        if (
            counters.invocations_per_task[task_name]
            > self.workflow.policies.max_invocations_per_task
        ):
            raise WorkflowLimitExceededError(
                f"Task {task_name} exceeded max_invocations_per_task"
            )

    def _enforce_run_timeout(self, run_started_at: float) -> None:
        timeout = self.workflow.policies.run_timeout_seconds
        if timeout is None:
            return
        if monotonic() - run_started_at > timeout:
            raise WorkflowTimeoutError(
                f"Workflow exceeded run_timeout_seconds ({timeout})"
            )
