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

from alayaflow.execution.executors.base_executor import BaseExecutor
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

_SENTINEL = object()

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, 
            input_data: dict, 
            user_config: dict
    ) -> Generator[Dict, None, None]:

        # 1) resolve workflow
        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

        print(f"NaiveExecutor execute_stream: {workflow_id} {version} {input_data} {user_config}")

        # TODO: Support langflow workflow
        # Only support StateGraphRunnableWorkflow now. 
        # As in _produce_events_to_queue, we relay on the "configurable" (StateGraph only) to use langfuse tracing.
        if not isinstance(runnable, StateGraphRunnableWorkflow):
            raise ValueError(f"NaiveExecutor only supports StateGraphRunnableWorkflow, but got {type(runnable)}")

        # 3) async -> sync bridge
        event_queue = queue.Queue(maxsize=1000)

        def run_async_producer():
            try:
                asyncio.run(self._produce_events_to_queue(runnable, input_data, user_config, event_queue))
            except Exception as e:
                event_queue.put({"error": str(e), "traceback": traceback.format_exc(), "workflow_id": workflow_id, "version": version})
            finally:
                event_queue.put(_SENTINEL)

        producer_thread = threading.Thread(target=run_async_producer, daemon=True)
        producer_thread.start()

        # 4) stream
        while True:
            try:
                item = event_queue.get(timeout=1.0)
            except queue.Empty:
                if not producer_thread.is_alive():
                    saw_sentinel = False
                    while True:
                        try:
                            item = event_queue.get_nowait()
                            if item is _SENTINEL:
                                saw_sentinel = True
                                break
                            yield self._serialize_event(item)
                        except queue.Empty:
                            break
                    if saw_sentinel:
                        break
                    yield {"error": "producer thread exited unexpectedly", "workflow_id": workflow_id, "version": version}
                    break
                continue
            if item is _SENTINEL:
                break
            yield self._serialize_event(item)

    
    async def _produce_events_to_queue(self, runnable: BaseRunnableWorkflow, input_data: dict, user_config: dict, event_queue: queue.Queue):
        try:
            # Setup tracing
            tracing = get_tracing(settings)
            langfuse_cb = tracing.build_callback()
            
            # Merge user_config and tracing config
            merged_config = {
                "configurable": user_config
            }
            if langfuse_cb:
                merged_config.update(tracing.build_config(input_data, runnable.info, langfuse_cb))
            
            async for chunk in runnable.astream_events(input_data, merged_config):
                event_queue.put(chunk)  # Put each event immediately (real-time)
        except Exception as e:
            # If execution fails, put error event in queue
            event_queue.put({
                "error": str(e),
                "traceback": traceback.format_exc(),
                "workflow_id": runnable.info.id,
                "version": runnable.info.version
            })

