import asyncio
import traceback
import queue
import threading
from typing import Generator, Dict, Any, AsyncGenerator

from alayaflow.execution.executors.base_executor import BaseExecutor
from alayaflow.utils.coroutine import iter_sync
from alayaflow.workflow.workflow_manager import WorkflowManager
from alayaflow.workflow.runnable import BaseRunnableWorkflow, StateGraphRunnableWorkflow
from alayaflow.common.config import settings
from alayaflow.execution.langfuse_tracing import get_tracing
from alayaflow.utils.logger import AlayaFlowLogger

logger = AlayaFlowLogger()


class NaiveExecutor(BaseExecutor):
    def __init__(self, workflow_manager: WorkflowManager):
        self.workflow_manager = workflow_manager

    def _to_jsonable(self, obj: Any) -> Any:
        """Recursively convert object to JSON-serializable format."""
        if obj is None or isinstance(obj, (bool, int, float, str)):
            return obj
        if isinstance(obj, dict):
            return {k: self._to_jsonable(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple)):
            return [self._to_jsonable(item) for item in obj]
        if hasattr(obj, 'model_dump') and callable(obj.model_dump):
            return obj.model_dump()
        return str(obj)
    
    def _serialize_event(self, event: Dict) -> Dict:
        # Convert event to JSON-serializable format
        return self._to_jsonable(event)

    def execute_stream(
            self, 
            workflow_id: str, 
            version: str, 
            inputs: dict, 
            context: dict,
            metadata: dict = None
    ) -> Generator[Dict, None, None]:
        return iter_sync(self.execute_stream_async(workflow_id, version, inputs, context, metadata))

    async def execute_stream_async(
            self, 
            workflow_id: str, 
            version: str, 
            inputs: dict, 
            context: dict,
            metadata: dict = None
    ) -> AsyncGenerator[Dict, None]:
        try:
            runnable = self.workflow_manager.get_workflow_runnable(workflow_id, version)
        except ValueError as e:
            yield {"error": str(e), "workflow_id": workflow_id, "version": version}
            return

        logger.info(f"Execute workflow.\n{workflow_id=} {version=}\n{inputs=}\n{context=}")

        # TODO: Support langflow workflow
        # Only support StateGraphRunnableWorkflow for now
        # check it as we pass a langgraph specific config here
        if not isinstance(runnable, StateGraphRunnableWorkflow):
            raise ValueError(f"NaiveExecutor only supports StateGraphRunnableWorkflow, but got {type(runnable)}")

        try:
            tracing = get_tracing(settings)
            langfuse_cb = tracing.build_callback()
        
            if langfuse_cb:
                config = tracing.build_config(inputs, runnable.info, langfuse_cb)
            else:
                config = {}

            if metadata and isinstance(metadata, dict):
                if 'metadata' not in config:
                    config['metadata'] = {}
                config['metadata'].update(metadata)
            
            async for chunk in runnable.stream_events_async(inputs, context, config):
                yield self._serialize_event(chunk)
        except Exception as e:
            yield {"error": str(e), "traceback": traceback.format_exc(), "workflow_id": workflow_id, "version": version}

