from __future__ import annotations
import traceback
import asyncio
import uuid
from copy import deepcopy
from typing import Any, AsyncIterator, Awaitable, Callable

from loguru import logger
from opentelemetry import context as otel_context_api
from opentelemetry import trace
from opentelemetry.trace import SpanKind

from ..db.database import DatabaseManager
from ..execution_context import (
    ExecutionContext,
    get_current_context,
    reset_current_context,
    set_current_context,
)
from .edge import Edge
from .node import Node
from .port import Port
from .trigger import Trigger
from ..db.database import DatabaseManager
from ..utils.workflow_clone import workflow_clone
from .workflow_callback import WorkflowCallback, BuiltinWorkflowCallback, CallbackEvent
from .workflow_config import WorkflowConfig
from .context import Context

class Workflow:
    def __init__(
        self,
        id: uuid.UUID,
        config: WorkflowConfig,
        nodes: list[Node],
        input_ports: list[Port],
        output_ports: list[Port],
        _handle_stream_output: Callable[[str, AsyncIterator[str]], Awaitable[None]] = None, # deprecated
        _handle_query_user: Callable[[str, str], Awaitable[str]] = None,
        database_manager: DatabaseManager = None,
        max_concurrent_runs: int = 10,
        callbacks: list[WorkflowCallback] = [],
        debug_version: bool = False,    # 是否为debug过程中的临时版本

        # for run
        task_id: uuid.UUID = None,
        real_trigger_node: Trigger = None,

        # global variables
        global_context: Context = None,
    ) -> None:
        self.id = id
        self.config = config
        self.nodes = nodes
        self.ready_nodes: list[Node] = []
        self.input_ports = input_ports
        self.output_ports = output_ports
        self._handle_stream_output = _handle_stream_output
        self._handle_query_user = _handle_query_user
        self.after_trigger_workflow = None
        self.result_port = Port("result", Any)
        self.database_manager = database_manager
        self.max_concurrent_runs = max_concurrent_runs
        self.run_semaphore = asyncio.Semaphore(max_concurrent_runs)
        self.callbacks = callbacks
        self.debug_version = debug_version
        self.task_id = task_id
        self.real_trigger_node = real_trigger_node
        self.global_context = global_context
        self._validate()
        self._tracer = trace.get_tracer("service_forge.workflow")
    
    @property
    def name(self) -> str:
        return self.config.name
    
    @property
    def version(self) -> str:
        return self.config.version
    
    @property
    def description(self) -> str:
        return self.config.description

    def register_callback(self, callback: WorkflowCallback) -> None:
        self.callbacks.append(callback)

    def unregister_callback(self, callback: WorkflowCallback) -> None:
        self.callbacks.remove(callback)

    async def call_callbacks(self, callback_type: CallbackEvent, *args, **kwargs) -> None:
        for callback in self.callbacks:
            if callback_type == CallbackEvent.ON_WORKFLOW_START:
                await callback.on_workflow_start(*args, **kwargs)
            elif callback_type == CallbackEvent.ON_WORKFLOW_END:
                await callback.on_workflow_end(*args, **kwargs)
            elif callback_type == CallbackEvent.ON_WORKFLOW_ERROR:
                await callback.on_workflow_error(*args, **kwargs)
            elif callback_type == CallbackEvent.ON_NODE_START:
                await callback.on_node_start(*args, **kwargs)
            elif callback_type == CallbackEvent.ON_NODE_END:
                await callback.on_node_end(*args, **kwargs)
            elif callback_type == CallbackEvent.ON_NODE_STREAM_OUTPUT:
                await callback.on_node_stream_output(*args, **kwargs)

    def add_nodes(self, nodes: list[Node]) -> None:
        for node in nodes:
            node.workflow = self
        self.nodes.extend(nodes)

    def remove_nodes(self, nodes: list[Node]) -> None:
        for node in nodes:
            self.nodes.remove(node)

    def load_config(self) -> None: ...

    def _validate(self) -> None:
        # DAG
        ...

    def get_input_port_by_name(self, name: str) -> Port:
        for input_port in self.input_ports:
            if input_port.name == name:
                return input_port
        return None

    def get_output_port_by_name(self, name: str) -> Port:
        for output_port in self.output_ports:
            if output_port.name == name:
                return output_port
        return None

    def get_trigger_node(self) -> Trigger:
        trigger_nodes = [node for node in self.nodes if isinstance(node, Trigger)]
        if not trigger_nodes:
            raise ValueError("No trigger nodes found in workflow.")
        if len(trigger_nodes) > 1:
            raise ValueError("Multiple trigger nodes found in workflow.")
        return trigger_nodes[0]

    async def _run_node_with_callbacks(self, node: Node) -> bool:
        await self.call_callbacks(CallbackEvent.ON_NODE_START, node=node)
        
        try:
            result = node.run()
            if hasattr(result, '__anext__'):
                await self.handle_node_stream_output(node, result)
            elif asyncio.iscoroutine(result):
                await result
        except Exception as e:
            await self.call_callbacks(CallbackEvent.ON_WORKFLOW_ERROR, workflow=self, node=node, error=e)
            logger.error(f"Error when running node {node.name}: {str(e)}, task_id: {self.task_id}")
            return False
        finally:
            await self.call_callbacks(CallbackEvent.ON_NODE_END, node=node)
        return True

    async def run_after_trigger(self) -> Any:
        logger.info(f"Running workflow: {self.name}")

        await self.call_callbacks(CallbackEvent.ON_WORKFLOW_START, workflow=self)

        self.ready_nodes = []
        for edge in self.get_trigger_node().output_edges:
            edge.end_port.trigger()

        for input_port in self.input_ports:
            if input_port.value is not None:
                input_port.port.node.fill_input(input_port.port, input_port.value)

        for node in self.nodes:
            for key in node.AUTO_FILL_INPUT_PORTS:
                if key[0] not in [edge.end_port.name for edge in node.input_edges]:
                    node.fill_input_by_name(key[0], key[1])

        while self.ready_nodes:
            nodes = self.ready_nodes.copy()
            self.ready_nodes = []

            tasks = []
            for node in nodes:
                tasks.append(asyncio.create_task(self._run_node_with_callbacks(node)))

            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    for task in tasks:
                        if not task.done():
                            task.cancel()
                    await asyncio.gather(*tasks, return_exceptions=True)
                    return
                    # raise result
                elif result is False:
                    logger.error(f"Node execution failed, stopping workflow: {nodes[i].name}")
                    for task in tasks:
                        if not task.done():
                            task.cancel()
                    await asyncio.gather(*tasks, return_exceptions=True)
                    return
                    # raise RuntimeError(f"Workflow stopped due to node execution failure: {nodes[i].name}")

        if len(self.output_ports) > 0:
            if len(self.output_ports) == 1:
                if self.output_ports[0].is_prepared:
                    result = self.output_ports[0].value
                else:
                    result = None
            else:
                result = {}
                for port in self.output_ports:
                    if port.is_prepared:
                        result[port.name] = port.value
            await self.call_callbacks(CallbackEvent.ON_WORKFLOW_END, workflow=self, output=result)
        else:
            await self.call_callbacks(CallbackEvent.ON_WORKFLOW_END, workflow=self, output=None)

    
    async def _run(self, task_id: uuid.UUID, trigger_node: Trigger) -> None:
        async with self.run_semaphore:
            base_context = get_current_context()

            # 尝试从 trigger_node 上取父 trace context（如果存在）
            trigger_parent_context = None
            if hasattr(trigger_node, "task_contexts"):
                trigger_parent_context = trigger_node.task_contexts.pop(task_id, None)

            parent_context = (
                trigger_parent_context
                or (
                    base_context.trace_context
                    if base_context and base_context.trace_context
                    else otel_context_api.get_current()
                )
            )

            span_name = f"Workflow {self.name}"
            token = None

            try:
                with self._tracer.start_as_current_span(
                    span_name,
                    context=parent_context,
                    kind=SpanKind.INTERNAL,
                ) as span:
                    span.set_attribute("workflow.name", self.name)
                    span.set_attribute("workflow.task_id", str(task_id))
                    span.set_attribute("service.name", "service_forge")

                    execution_context = ExecutionContext(
                        trace_context=otel_context_api.get_current(),
                        span=span,
                        metadata={
                            **(base_context.metadata if base_context else {}),
                            "workflow_name": self.name,
                            "task_id": str(task_id),
                        },
                    )
                    token = set_current_context(execution_context)

                    new_workflow = self._clone(task_id, trigger_node)

                    # run_after_trigger 当前实现没有 return，result 会是 None
                    result = await new_workflow.run_after_trigger()

                    # 将结果回传给等待方（如果有队列）
                    if hasattr(trigger_node, "result_queues") and task_id in trigger_node.result_queues:
                        trigger_node.result_queues[task_id].put_nowait(result)

                    # TODO: clear new_workflow
                    for node in new_workflow.nodes:
                        await node.clear()

            except Exception as e:
                await self.call_callbacks(CallbackEvent.ON_WORKFLOW_ERROR, workflow=self, node=None, error=e)
                # error_msg = f"Error running workflow: {str(e)}, {traceback.format_exc()}"
                # logger.error(error_msg)
                # await self.call_callbacks(CallbackEvent.ON_WORKFLOW_END, workflow=self, node=None, error=e)
                return

                # 更新任务状态为失败 + WebSocket 通知
                websocket_manager.task_manager.fail_task(task_id, error_msg)
                await websocket_manager.send_execution_error(task_id, "workflow", error_msg)

                # 异常也回传给等待方，避免对方一直 await
                if hasattr(trigger_node, "result_queues") and task_id in trigger_node.result_queues:
                    trigger_node.result_queues[task_id].put_nowait(e)

                # 将异常继续抛出，便于上层感知失败
                raise
            finally:
                if token is not None:
                    reset_current_context(token)

    async def run(self):
        tasks = []
        trigger = self.get_trigger_node()

        async for task_id in trigger.run():
            tasks.append(asyncio.create_task(self._run(task_id, trigger)))

        if tasks:
            await asyncio.gather(*tasks)

    async def stop(self):
        trigger = self.get_trigger_node()
        await trigger._stop()

    def trigger(self, trigger_name: str, assigned_task_id: uuid.UUID | None, **kwargs) -> uuid.UUID:
        trigger = self.get_trigger_node()
        task_id = assigned_task_id or uuid.uuid4()
        for key, value in kwargs.items():
            trigger.prepare_output_edges(key, value)
        task = asyncio.create_task(self._run(task_id, trigger))
        return task_id

    async def handle_node_stream_output(
        self,
        node: Node,
        stream: AsyncIterator[Any],
    ) -> None:
        async for data in stream:
            await self.call_callbacks(CallbackEvent.ON_NODE_STREAM_OUTPUT, node=node, output=data)

    # TODO: refactor this
    async def handle_query_user(self, node_name: str, prompt: str) -> Awaitable[str]:
        return await asyncio.to_thread(input, f"[{node_name}] {prompt}: ")
    
    def _clone(self, task_id: uuid.UUID, trigger_node: Trigger) -> Workflow:
        return workflow_clone(self, task_id, trigger_node)
