
from __future__ import annotations
import asyncio
import uuid
from typing import AsyncIterator, Awaitable, Callable, Any
from loguru import logger
from copy import deepcopy
from .node import Node
from .port import Port
from .trigger import Trigger
from .edge import Edge
from ..db.database import DatabaseManager
from ..api.websocket_manager import websocket_manager


class Workflow:
    def __init__(
        self,
        name: str,
        description: str,
        nodes: list[Node],
        input_ports: list[Port],
        output_ports: list[Port],
        _handle_stream_output: Callable[[str, AsyncIterator[str]], Awaitable[None]] = None,
        _handle_query_user: Callable[[str, str], Awaitable[str]] = None,
        database_manager: DatabaseManager = None,
        max_concurrent_runs: int = 10,
    ) -> None:
        self.name = name
        self.description = description
        self.nodes = nodes
        self.ready_nodes = []
        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.run_semaphore = asyncio.Semaphore(max_concurrent_runs)
        self._validate()

    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

    async def run_after_trigger(self, task_id: uuid.UUID = None) -> Any:
        trigger_nodes = [node for node in self.nodes if isinstance(node, Trigger)]
        if not trigger_nodes:
            raise ValueError("No trigger nodes found in workflow.")
        self.ready_nodes = []
        for edge in trigger_nodes[0].output_edges:
            edge.end_port.trigger()

        logger.info(f"Running workflow: {self.name}")
        result = None
        try:
            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:
                tasks = []
                nodes = self.ready_nodes.copy()
                self.ready_nodes = []
                for i, node in enumerate(nodes):
                    # 发送节点开始执行的WebSocket通知
                    if task_id:
                        # 更新当前步骤
                        websocket_manager.task_manager.update_current_step(task_id, i + 1)
                        await websocket_manager.send_executing(task_id, node.name)

                    result = node.run(task_id)

                    if hasattr(result, '__anext__'):
                        # 处理流输出
                        if self._handle_stream_output is None:
                            tasks.append(self.handle_stream_output(node.name, result))
                        else:
                            tasks.append(self._handle_stream_output(node.name, result))
                    elif asyncio.iscoroutine(result):
                        tasks.append(result)

                    # 发送每个输出端口的结果
                    if task_id:
                        for output_port in node.output_ports:
                            if output_port.is_prepared:
                                await websocket_manager.send_node_output(task_id, node.name, output_port.name, output_port.value)
                
                if tasks:
                    results = await asyncio.gather(*tasks)
                    # 更新result为最后一个非None的结果
                    for res in results:
                        if res is not None:
                            result = res

                    # 在所有任务完成后，再次发送每个输出端口的结果
                    if task_id:
                        for node in nodes:
                            for output_port in node.output_ports:
                                if output_port.is_prepared:
                                    await websocket_manager.send_node_output(task_id, node.name, output_port.name, output_port.value)

        except Exception as e:
            import traceback
            error_msg = f"Error in run_after_trigger: {str(e)}"
            logger.error(error_msg)
            logger.error(traceback.format_exc())
            # 发送执行错误的WebSocket通知
            if task_id:
                # 更新任务状态为失败
                websocket_manager.task_manager.fail_task(task_id, error_msg)
                await websocket_manager.send_execution_error(task_id, "workflow", error_msg)
            raise e

        # 发送工作流执行完成的WebSocket通知
        if task_id:
            # 更新任务状态为已完成
            websocket_manager.task_manager.complete_task(task_id)
            await websocket_manager.send_executed(task_id, "workflow", result)

        if len(self.output_ports) > 0 and self.output_ports[0].is_prepared:
            return self.output_ports[0].value

    async def _run(self, uuid: uuid.UUID, trigger_node: Trigger) -> None:
        async with self.run_semaphore:
            try:
                # 发送任务开始执行的WebSocket通知
                await websocket_manager.send_execution_start(uuid)

                new_workflow = self._clone()
                result = await new_workflow.run_after_trigger(uuid)
                if uuid in trigger_node.result_queues:
                    trigger_node.result_queues[uuid].put_nowait(result)

            except Exception as e:
                import traceback
                error_msg = f"Error running workflow: {str(e)}, {traceback.format_exc()}"
                logger.error(error_msg)
                # 发送执行错误的WebSocket通知
                # 更新任务状态为失败
                websocket_manager.task_manager.fail_task(uuid, error_msg)
                await websocket_manager.send_execution_error(uuid, "workflow", error_msg)
                if uuid in trigger_node.result_queues:
                    trigger_node.result_queues[uuid].put_nowait(e)

    async def run(self):
        trigger_nodes = [node for node in self.nodes if isinstance(node, Trigger)]

        if len(trigger_nodes) == 0:
            raise ValueError("No trigger nodes found in workflow.")
        if len(trigger_nodes) > 1:
            raise ValueError("Multiple trigger nodes found in workflow.")

        trigger_node = trigger_nodes[0]
        tasks = []
        async for uuid in trigger_node.run():
            tasks.append(asyncio.create_task(self._run(uuid, trigger_node)))
        
        if tasks:
            await asyncio.gather(*tasks)

    async def handle_stream_output(self, node_name: str, stream: AsyncIterator[str]) -> None:
        logger.info(f"[{node_name}] Starting stream output:")
        buffer = []
        async for char in stream:
            buffer.append(char)
            logger.info(f"[{node_name}] Received char: '{char}'")

        complete_message = ''.join(buffer)
        logger.info(f"[{node_name}] Complete message: '{complete_message}'")

    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) -> Workflow:
        node_map: dict[Node, Node] = {node: node._simple_clone() for node in self.nodes}

        port_map: dict[Port, Port] = {}
        port_map.update({port: port._simple_clone(node_map) for port in self.input_ports})
        port_map.update({port: port._simple_clone(node_map) for port in self.output_ports})
        for node in self.nodes:
            for port in node.input_ports:
                if port not in port_map:
                    port_map[port] = port._simple_clone(node_map)
            for port in node.output_ports:
                if port not in port_map:
                    port_map[port] = port._simple_clone(node_map)

        edge_map: dict[Edge, Edge] = {}
        for node in self.nodes:
            for edge in node.input_edges:
                if edge not in edge_map:
                    edge_map[edge] = edge._simple_clone(node_map, port_map)
            for edge in node.output_edges:
                if edge not in edge_map:
                    edge_map[edge] = edge._simple_clone(node_map, port_map)

        # fill port.port
        for old_port, new_port in port_map.items():
            if old_port.port is not None:
                new_port.port = port_map[old_port.port]


        # fill ports and edges in nodes
        for old_node, new_node in node_map.items():
            new_node.input_edges = [edge_map[edge] for edge in old_node.input_edges]
            new_node.output_edges = [edge_map[edge] for edge in old_node.output_edges]
            new_node.input_ports = [port_map[port] for port in old_node.input_ports]
            new_node.output_ports = [port_map[port] for port in old_node.output_ports]
            new_node.input_variables = {
                port_map[port]: value for port, value in old_node.input_variables.items()
            }

        workflow = Workflow(
            name=self.name,
            description=self.description,
            nodes=[node_map[node] for node in self.nodes],
            input_ports=[port_map[port] for port in self.input_ports],
            output_ports=[port_map[port] for port in self.output_ports],
            _handle_stream_output=self._handle_stream_output,
            _handle_query_user=self._handle_query_user,
            database_manager=self.database_manager,
        )

        for node in workflow.nodes:
            node.workflow = workflow

        return workflow
