from bedrock_agentcore.runtime import BedrockAgentCoreApp
from llama_agents.client import HandlerData
from llama_agents.server._runtime.idle_release_runtime import IdleReleaseDecorator
from llama_agents.server._runtime.persistence_runtime import PersistenceDecorator
from llama_agents.server._service import _WorkflowService as WorkflowService
from llama_agents.server._store.abstract_workflow_store import AbstractWorkflowStore
from workflows import Workflow
from workflows.events import StartEvent
from workflows.plugins.basic import basic_runtime
from workflows.runtime.types.plugin import Runtime
from workflows.utils import _nanoid

from ._runtime_decorator import AgentCoreRuntimeDecorator

IDLE_TIMEOUT = 60.0
PERSISTENCE_BACKOFF = [0.5, 3]
POLLING_TIMEOUT = 600


class WorkflowNotFoundError(Exception):
    """
    Raise when a workflow with a given name cannot be found within the registered workflows.
    """

    def __init__(self, workflow_name: str) -> None:
        self.workflow_name = workflow_name

    def __repr__(self) -> str:
        return f"Could not find {self.workflow_name} among the registered workflows"

    def __str__(self) -> str:
        return f"Could not find {self.workflow_name} among the registered workflows"


class AgentCoreService:
    def __init__(
        self,
        app: BedrockAgentCoreApp,
        store: AbstractWorkflowStore,
        idle_timeout: float = IDLE_TIMEOUT,
        persistence_backoff: list[float] = PERSISTENCE_BACKOFF,
    ) -> None:
        self._workflow_store = store
        inner: Runtime = IdleReleaseDecorator(
            PersistenceDecorator(basic_runtime, store=self._workflow_store),
            store=self._workflow_store,
            idle_timeout=idle_timeout,
        )
        self._runtime: AgentCoreRuntimeDecorator = AgentCoreRuntimeDecorator(
            decorated=inner,
            store=self._workflow_store,
            app=app,
            persistence_backoff=persistence_backoff,
        )
        self._service = WorkflowService(
            runtime=self._runtime, store=self._workflow_store
        )

    def add_workflow(self, workflow_name: str, workflow: Workflow) -> None:
        workflow._switch_workflow_name(workflow_name)
        workflow._switch_runtime(self._runtime)

    async def run_workflow(
        self, workflow_name: str, start_event: StartEvent
    ) -> HandlerData:
        wf = self._service.get_workflow(workflow_name)
        if wf is not None:
            handler_data = await self._service.start_workflow(
                wf, start_event=start_event, handler_id=_nanoid()
            )
            return await self._service.await_workflow(handler_data)

        raise WorkflowNotFoundError(workflow_name)
