from typing import Dict, Generator, Any, AsyncGenerator
from enum import Enum

from alayaflow.execution.executors.base_executor import BaseExecutor
from alayaflow.execution.executors.uv_executor import UvExecutor
from alayaflow.execution.executors.naive_executor import NaiveExecutor
from alayaflow.execution.executors.worker_executor import WorkerExecutor
from alayaflow.workflow import WorkflowManager


class ExecutorType(Enum):
    UV = "uv"
    NAIVE = "naive"
    WORKER = "worker"


class ExecutorManager:    
    def __init__(
        self,
        workflow_manager: WorkflowManager
    ):
        self._workflow_manager = workflow_manager
        self._executor_map: Dict[ExecutorType, BaseExecutor] = {}
        self._initialize_executors()
    
    def _initialize_executors(self):
        """Initialize all available executors."""
        self._executor_map[ExecutorType.UV] = UvExecutor(
            workflow_manager=self._workflow_manager
        )
        self._executor_map[ExecutorType.NAIVE] = NaiveExecutor(
            workflow_manager=self._workflow_manager
        )
        self._executor_map[ExecutorType.WORKER] = WorkerExecutor(
            workflow_manager=self._workflow_manager
        )
        
    def init_workflow(self, workflow_id: str, version: str, init_args: dict):
        for executor in self._executor_map.values():
            executor.init_workflow(workflow_id, version, init_args)
    
    def exec_workflow(
            self, 
            workflow_id: str,
            version: str, 
            inputs: dict,
            context: dict, 
            executor_type: ExecutorType | str = ExecutorType.NAIVE,
            metadata: dict = None
    ) -> Generator[Dict[str, Any], None, None]:
        if isinstance(executor_type, str):
            executor_type = ExecutorType(executor_type)
        if executor_type not in self._executor_map:
            raise ValueError(
                f"Unsupported executor kind: {executor_type}. "
                f"Supported kinds: {list(self._executor_map.keys())}"
            )
        executor = self._executor_map[executor_type]
        yield from executor.execute_stream(workflow_id, version, inputs, context, metadata)

    async def exec_workflow_async(
            self, 
            workflow_id: str,
            version: str, 
            inputs: dict,
            context: dict, 
            executor_type: ExecutorType | str = ExecutorType.NAIVE,
            metadata: dict = None
    ) -> AsyncGenerator[Dict[str, Any], None]:
        if isinstance(executor_type, str):
            executor_type = ExecutorType(executor_type)
        if executor_type not in self._executor_map:
            raise ValueError(
                f"Unsupported executor kind: {executor_type}. "
                f"Supported kinds: {list(self._executor_map.keys())}"
            )
        executor = self._executor_map[executor_type]
        async for event in executor.execute_stream_async(workflow_id, version, inputs, context, metadata):
            yield event