from asyncio import AbstractEventLoop, CancelledError, Task, TimerHandle, gather, get_running_loop
from inspect import isfunction
from typing import Protocol, cast, final, overload, runtime_checkable
from uuid import UUID, uuid4

from draive.agents.definition import AgentMessage, AgentNode
from draive.agents.errors import AgentException
from draive.agents.idle import IdleMonitor
from draive.agents.node import Agent, AgentOutput
from draive.agents.runner import AgentRunner
from draive.helpers import VolatileMemory
from draive.parameters import DataModel, ParametrizedData, State
from draive.scope import ctx
from draive.types import (
    BasicMemory,
    Memory,
    MultimodalContent,
    MultimodalContentConvertible,
    frozenlist,
)
from draive.utils import MISSING, AsyncQueue, Missing, freeze, not_missing

__all__ = [
    "workflow",
    "AgentWorkflow",
    "AgentWorkflowInvocation",
    "AgentWorkflowOutput",
    "AgentWorkflowIdle",
    "AgentWorkflowInput",
]


type AgentWorkflowInput = AgentMessage | AgentException
type AgentWorkflowOutput[Result: ParametrizedData | str] = (
    frozenlist[AgentMessage] | AgentMessage | Result | None
)


@runtime_checkable
class AgentWorkflowStateInitializer[AgentWorkflowState](Protocol):
    def __call__(self) -> AgentWorkflowState: ...


@runtime_checkable
class AgentWorkflowInvocation[AgentWorkflowState, AgentWorkflowResult: ParametrizedData | str](
    Protocol
):
    async def __call__(
        self,
        memory: BasicMemory[AgentWorkflowState],
        input: AgentWorkflowInput,  # noqa: A002
    ) -> AgentWorkflowOutput[AgentWorkflowResult]: ...


class AgentWorkflowIdle(AgentException):
    pass


@final
class AgentWorkflow[AgentWorkflowState, AgentWorkflowResult: ParametrizedData | str]:
    def __init__(
        self,
        node: AgentNode,
        invocation: AgentWorkflowInvocation[
            AgentWorkflowState,
            AgentWorkflowResult,
        ],
        state_initializer: AgentWorkflowStateInitializer[AgentWorkflowState],
    ) -> None:
        self.node: AgentNode = node
        self._invocation: AgentWorkflowInvocation[
            AgentWorkflowState,
            AgentWorkflowResult,
        ] = invocation
        self._state_initializer: AgentWorkflowStateInitializer[AgentWorkflowState] = (
            state_initializer
        )

        freeze(self)

    async def __call__(
        self,
        memory: BasicMemory[AgentWorkflowState],
        input: AgentWorkflowInput,  # noqa: A002
    ) -> AgentWorkflowOutput[AgentWorkflowResult]:
        return await self._invocation(
            memory=memory,
            input=input,
        )

    async def run(
        self,
        input: MultimodalContent | MultimodalContentConvertible,  # noqa: A002
        state: AgentWorkflowState | None = None,
        timeout: float = 120,  # default timeout is 2 minutes
    ) -> AgentWorkflowResult:
        return await WorkflowRunner[AgentWorkflowState, AgentWorkflowResult].run(
            self,
            input=input,
            memory=VolatileMemory(state or self._state_initializer()),
            timeout=timeout,
        )

    def address(
        self,
        content: MultimodalContent | MultimodalContentConvertible,
        /,
        *_content: MultimodalContent | MultimodalContentConvertible,
        addressee: AgentNode | None = None,
    ) -> AgentMessage:
        return AgentMessage(
            identifier=uuid4(),
            sender=MISSING,
            recipient=self.node,
            addressee=addressee,
            content=MultimodalContent.of(content, *_content),
            responding=None,
        )


class PartialAgentWorkflowWrapper[AgentWorkflowState](Protocol):
    def __call__[AgentWorkflowResult: ParametrizedData | str](
        self,
        invocation: AgentWorkflowInvocation[AgentWorkflowState, AgentWorkflowResult],
    ) -> AgentWorkflow[AgentWorkflowState, AgentWorkflowResult]: ...


@overload
def workflow[AgentWorkflowState](
    node: AgentNode,
    /,
    *,
    state: AgentWorkflowStateInitializer[AgentWorkflowState],
) -> PartialAgentWorkflowWrapper[AgentWorkflowState]: ...


@overload
def workflow[AgentWorkflowState](
    *,
    name: str | None = None,
    description: str,
    state: AgentWorkflowStateInitializer[AgentWorkflowState],
) -> PartialAgentWorkflowWrapper[AgentWorkflowState]: ...


def workflow[AgentWorkflowState](
    node: AgentNode | None = None,
    *,
    name: str | None = None,
    description: str | None = None,
    state: AgentWorkflowStateInitializer[AgentWorkflowState],
) -> PartialAgentWorkflowWrapper[AgentWorkflowState]:
    assert node is None or (  # nosec: B101
        name is None and description is None
    ), "Can't specify both agent node and name/description"
    assert (  # nosec: B101
        description is not None or node is not None
    ), "Either agent node or description has to be provided"

    def wrap[AgentWorkflowResult: ParametrizedData | str](
        invocation: AgentWorkflowInvocation[AgentWorkflowState, AgentWorkflowResult],
    ) -> AgentWorkflow[AgentWorkflowState, AgentWorkflowResult]:
        assert isfunction(invocation), "workflow has to be defined from a function"  # nosec: B101

        workflow_node: AgentNode = node or AgentNode(
            name=name or invocation.__qualname__,
            description=description or "",
        )

        def initialize_agent() -> Agent:
            agent_memory: BasicMemory[AgentWorkflowState] = cast(
                BasicMemory[AgentWorkflowState], VolatileMemory(state())
            )

            async def agent(message: AgentMessage) -> AgentOutput:
                match await cast(  # pyright fails to properly type this function
                    AgentWorkflowInvocation[AgentWorkflowState, AgentWorkflowResult],
                    invocation,
                )(
                    memory=agent_memory,
                    input=message,
                ):
                    case None:
                        pass

                    case tuple() as messages:
                        return messages

                    case AgentMessage() as message:
                        return message

                    case result:
                        return message.respond(
                            MultimodalContent.of(
                                result if isinstance(result, DataModel) else str(result)
                            )
                        )

            return agent

        # workflow can run as an agent with limited capabilities
        workflow_node.associate(
            initialize_agent,
            concurrent=False,  # workflow can't be concurrent
        )

        return AgentWorkflow[AgentWorkflowState, AgentWorkflowResult](
            node=workflow_node,
            invocation=invocation,
            state_initializer=state,
        )

    return wrap


@runtime_checkable
class WorkflowRunnerOutput(Protocol):
    def __call__(
        self,
        messages: AgentMessage,
        /,
        *_messages: AgentMessage,
    ) -> None: ...


@runtime_checkable
class WorkflowRunnerResultOutput[WorkflowResult: ParametrizedData | str](Protocol):
    def __call__(
        self,
        result: WorkflowResult,
        /,
    ) -> None: ...


class WorkflowHistory(State):
    history: frozenlist[AgentMessage]


class WorkflowRunner[WorkflowState, WorkflowResult: ParametrizedData | str]:
    @classmethod
    async def run(
        cls,
        workflow: AgentWorkflow[WorkflowState, WorkflowResult],
        /,
        input: MultimodalContent | MultimodalContentConvertible,  # noqa: A002
        memory: Memory[WorkflowState, WorkflowState],
        timeout: float,
    ) -> WorkflowResult:
        return await cls(
            workflow=workflow,
            memory=memory,
            loop=get_running_loop(),
        ).execute(
            input=input,
            timeout=timeout,
        )

    def __init__(
        self,
        workflow: AgentWorkflow[WorkflowState, WorkflowResult],
        memory: Memory[WorkflowState, WorkflowState],
        loop: AbstractEventLoop,
    ) -> None:
        self._loop: AbstractEventLoop = loop
        self._history: list[AgentMessage] = []
        self._workflow_queue: AsyncQueue[AgentMessage | AgentException] = AsyncQueue()
        self._workflow_memory: Memory[WorkflowState, WorkflowState] = memory
        self._workflow: AgentWorkflow[WorkflowState, WorkflowResult] = workflow
        self._agent_runners: dict[UUID, AgentRunner] = {}
        self._status: IdleMonitor = IdleMonitor()
        self._result: WorkflowResult | BaseException | Missing = MISSING

    def __del__(self) -> None:
        self.finish(result=CancelledError())

    @overload
    def send(
        self,
        error: AgentException,
        /,
    ) -> None: ...

    @overload
    def send(
        self,
        messages: AgentMessage,
        /,
        *_messages: AgentMessage,
    ) -> None: ...

    def send(
        self,
        messages: AgentMessage | AgentException,
        /,
        *_messages: AgentMessage,
    ) -> None:
        pending: list[AgentMessage | AgentException] = [messages, *_messages]
        for _ in pending:  # ensure not idle when received a message
            self._status.enter_task()
        self._history.extend([message for message in pending if isinstance(message, AgentMessage)])
        self._workflow_queue.enqueue(*pending)

    @property
    def finished(self) -> bool:
        return not_missing(self._result)

    def finish(
        self,
        result: WorkflowResult | BaseException,
    ) -> None:
        if not_missing(self._result):
            return  # already finished

        ctx.log_debug("Finishing workflow with result: %s", result)

        self._result = result

        for runner in self._agent_runners.values():
            runner.finish()

        self._workflow_queue.finish()

    async def _wait(self) -> None:
        await gather(
            *[runner.wait() for runner in self._agent_runners.values()],
            return_exceptions=False,
        )

    async def finalize(
        self,
        result: WorkflowResult | BaseException,
    ) -> None:
        self.finish(result=result)
        await self._wait()

    async def execute(  # noqa: PLR0912
        self,
        input: MultimodalContent | MultimodalContentConvertible,  # noqa: A002
        timeout: float,
    ) -> WorkflowResult:
        assert not self._agent_runners, "WorkflowRunner can run only once!"  # nosec: B101

        if self._workflow_queue.finished or not_missing(self._result):
            raise RuntimeError("WorkflowRunner can run only once!")

        async with ctx.nested(self._workflow.node.__str__()):

            def on_timeout() -> None:
                self.finish(result=TimeoutError())

            timeout_handle: TimerHandle = self._loop.call_later(
                delay=timeout,
                callback=on_timeout,
            )

            self.send(self._workflow.address(input))

            idle_monitor_task: Task[None] = ctx.spawn_subtask(self._idle_monitor)

            try:
                async for element in self._workflow_queue:
                    try:
                        match element:
                            case AgentMessage() as message:
                                if message.recipient.identifier == self._workflow.node.identifier:
                                    await self._handle(message)

                                elif runner := self._agent_runners.get(
                                    message.recipient.identifier
                                ):
                                    runner.send(message)

                                else:
                                    spawned_runner: AgentRunner = AgentRunner.run(
                                        message.recipient,
                                        output=self.send,
                                        status=self._status.nested(),
                                    )
                                    self._agent_runners[message.recipient.identifier] = (
                                        spawned_runner
                                    )
                                    spawned_runner.send(message)

                            case error:
                                await self._handle(error)

                    finally:
                        self._status.exit_task()

            except CancelledError:
                pass  # just end on cancel

            finally:  # cleanup status when exiting
                timeout_handle.cancel()  # cancel the timeout
                idle_monitor_task.cancel()  # finish idle monitor
                self._status.exit_all()

            await self._wait()  # wait for completion of all runners

            ctx.record(WorkflowHistory(history=tuple(self._history)))

            match self._result:
                case BaseException() as exc:
                    raise exc

                case Missing():
                    raise RuntimeError("Invalid workflow state")

                case result:
                    return result

    async def _handle(
        self,
        input: AgentWorkflowInput,  # noqa: A002
        /,
    ) -> None:
        try:
            match await self._workflow(
                memory=self._workflow_memory,
                input=input,
            ):
                case None:
                    pass  # nothing to do

                case [*messages]:
                    self.send(
                        *[message.updated(sender=self._workflow.node) for message in messages]
                    )

                case AgentMessage() as message:
                    self.send(message.updated(sender=self._workflow.node))

                case result:
                    self.finish(result=cast(WorkflowResult, result))

        except BaseException as exc:
            self.finish(result=exc)

    async def _idle_monitor(self) -> None:
        while not self.finished:  # run until finished
            await self._status.wait_idle()  # wait until all agents become idle

            if self.finished:  # check if that is not due to finishing
                return  # workflow might be finished at this stage

            ctx.log_warning("Detected workflow idle state, attempting to resume...")

            # when workflow should still be running (did not provided result yet)
            # but no agent is working then notify workflow node about the error
            self.send(AgentWorkflowIdle())
