from abc import ABC, abstractmethod
from typing import Iterator, List, Dict, Any, Set

# ===== DAG 基础结构 =====
# from arrow_flight_project.client.remote_dataframe import RemoteDataFrameImpl
# # from arrow_flight_project.client.arrow_flight_client import ArrowFlightClient
# from arrow_flight_project.client.operation import Operation, SelectOp
# from arrow_flight_project.udf.filters import filter_id_less_than_10
from arrow_flight_project.client.operation import Schema, SourceOp, FilterOp, MapOp, LimitOp, FunctionWrapper


class Row:
    def __init__(self, data: Dict[str, Any]):
        self.data = data


class DAGNode(ABC):
    pass


class UDFFunction(DAGNode):
    def __init__(self, func_wrapper: FunctionWrapper):
        self.func_wrapper = func_wrapper

    @abstractmethod
    def transform(self, rows: Iterator[Row]) -> Iterator[Row]:
        pass


class ConcreteUDFFunction(UDFFunction):
    def transform(self, rows: Iterator[Row]) -> Iterator[Row]:
        # This method is actually called by TransformerNode.execute,
        # where the FunctionWrapper's apply_to_input handles the dispatch.
        # This transform method itself might not be directly called in the current setup
        # if TransformerNode.execute directly uses func_wrapper.apply_to_input.
        # But keeping it for conceptual clarity of UDFFunction interface.
        return self.func_wrapper.apply_to_input(rows)


class SourceNode(DAGNode):
    def __init__(self, data_frame_name: str):
        self.data_frame_name = data_frame_name


# Added for the example extending DAG capabilities
class LimitNode(DAGNode):
    def __init__(self, limit_value: int):
        self.limit_value = limit_value


class SelectNode(DAGNode):
    def __init__(self, *columns: str):
        self.columns = list(columns)


# ===== Transformer DAG 实现 =====

class TransformerDAG:
    def __init__(self, nodes: Dict[str, DAGNode], edges: Dict[str, List[str]]):
        self.nodes = nodes
        self.edges = edges

    def get_execution_paths(self) -> List[List[DAGNode]]:
        key_paths = self._extract_all_paths(self.edges)
        node_paths = []
        for path in key_paths:
            node_path = []
            for key in path:
                node = self.nodes.get(key)
                if node is None:
                    raise ValueError(f"Invalid DAG: node '{key}' is not defined in the node map.")
                node_path.append(node)
            node_paths.append(node_path)
        return node_paths

    def _extract_all_paths(self, edges: Dict[str, List[str]]) -> List[List[str]]:
        all_targets = {target for children in edges.values() for target in children}
        all_sources = set(edges.keys())
        root_nodes = all_sources - all_targets

        if not root_nodes:
            raise ValueError("Invalid DAG: no root nodes found. Graph might contain cycles or be empty.")

        for root_key in root_nodes:
            node = self.nodes.get(root_key)
            if isinstance(node, SourceNode):
                continue
            elif node is None:
                raise ValueError(f"Invalid DAG: root node '{root_key}' is not defined in the node map.")
            else:
                raise ValueError(
                    f"Invalid DAG: root node '{root_key}' is not of type SourceNode, but {type(node).__name__}.")

        def dfs(path: List[str], current: str, visited: Set[str]) -> List[List[str]]:
            if current in visited:
                raise ValueError(f"Cycle detected: node '{current}' is revisited in path {' -> '.join(path)}")

            if current in edges:
                children = edges[current]
                all_paths = []
                for child in children:
                    all_paths.extend(dfs(path + [child], child, visited | {current}))
                return all_paths
            else:
                return [path]

        all_paths = []
        for root in root_nodes:
            all_paths.extend(dfs([root], root, set()))
        return all_paths


# def build_and_execute_dag_pipeline(
#         dag: TransformerDAG,
#         client: Any,
# ) -> List[Row]:
#     # 1. Validate and get execution paths
#     execution_paths = dag.get_execution_paths()
#
#     if not execution_paths:
#         raise ValueError("No execution paths found in the DAG.")
#
#     # Find the path that starts with the initial_data_frame_name's SourceNode
#     target_path = None
#
#     candidate_paths = [
#         path for path in execution_paths
#         # if isinstance(path[0], SourceNode) and path[0].data_frame_name == initial_data_frame_name
#     ]
#     if not candidate_paths:
#         raise ValueError(f"No execution path found.")
#
#     # 2. Build the chained Operation object
#     current_op: Operation = None
#     multiple_result_rows = []
#
#     for target_path in candidate_paths:
#         for i, node in enumerate(target_path):
#             if isinstance(node, SourceNode):
#                 if i != 0:
#                     raise ValueError("SourceNode must be the first node in an execution path.")
#                 current_op = SourceOp()  # Represents the starting point of the operation chain
#                 source_name = node.data_frame_name
#             elif isinstance(node, ConcreteUDFFunction):
#                 # Heuristics to map ConcreteUDFFunction to MapOp, FilterOp, or TransformerNode
#                 # based on the UDF's name. This is a simplification.
#                 func_name = node.func_wrapper.impl.__name__.lower()
#                 if "filter" in func_name:
#                     current_op = FilterOp(node.func_wrapper, current_op)
#                 elif "map" in func_name:
#                     current_op = MapOp(node.func_wrapper, current_op)
#                 else:
#                     raise TypeError(
#                         f"Unsupported func type")
#                 # else:
#                 #     # Default for generic UDFs that process streams
#                 #     current_op = TransformerNode(node.func_wrapper, current_op)
#             elif isinstance(node, LimitNode):  # Added support for LimitNode
#                 current_op = LimitOp(node.limit_value, current_op)
#             # Add more `elif` blocks here for other custom DAGNode types
#             elif isinstance(node, SelectNode):
#                 current_op = SelectOp(current_op, *node.columns)
#             else:
#                 raise TypeError(
#                     f"Unsupported DAG node type: {type(node).__name__}. Please ensure it's mapped to an Operation.")
#         if current_op is None:
#             raise ValueError("Failed to build any operation from the DAG path.")
#         # 3. Execute the chained operation
#         initial_df_for_execution = RemoteDataFrameImpl(source_name, client, None)
#         # Execute the final chained operation
#         result_rows = initial_df_for_execution.execute(current_op)
#         multiple_result_rows.append(result_rows)
#         # print(multiple_result_rows)
#     return multiple_result_rows


# ===== ✅ 测试用例 =====
# if __name__ == "__main__":
#
#     client = ArrowFlightClient("localhost", 3101)
#     client.login({"username": "admin@instdb.cn", "password": "admin001"})
#     dataset = client.list_data_set_names()
#     print(dataset)
#
#     #    adminUsername -> Set(s"$csvDir\\data_1.csv", s"$csvDir\\invalid.csv", s"$binDir", s"/csv/data_1.csv", "/bin",
#     #    "/csv/data_2.csv", "/bin/data_1.csv", "/csv/invalid.csv", "csv")
#     df = RemoteDataFrameImpl("csv", client)
#
#     # 模拟节点图结构如下：
#     source_node_1 = SourceNode("csv")
#     source_node_2 = SourceNode('/csv/data_1.csv')
#
#     # Functions for MapOp/FilterOp (process single Row)
#     filter_func_wrapper = FunctionWrapper(filter_id_less_than_10)
#     map_func_wrapper = FunctionWrapper(map_value_squre_10)
#
#     # New LimitNode for DAG structure
#     limit_node_2 = LimitNode(2)
#     limit_node_3 = LimitNode(3)
#
#     select_node_value = SelectNode("value")
#
#     dag_nodes = {
#         "source_1": source_node_1,
#         "source_2": source_node_2,
#         "map_op": ConcreteUDFFunction(map_func_wrapper),  # Maps value * 10
#         "filter_op": ConcreteUDFFunction(filter_func_wrapper),  # Filters id < 10
#         "limit_2_rows": limit_node_2,
#         "limit_3_rows": limit_node_3,
#         "select_value": select_node_value
#     }
#     linear_edges = {
#         # "source_1": ["map_op", "filter_op"],
#         "source_1": ["map_op"],
#         "source_2": ["map_op"],
#         "map_op": ["filter_op"],
#         "filter_op": ["limit_2_rows", "limit_3_rows"],
#         # "limit_2_rows": ["select_value"],
#         # "limit_3_rows": ["limit_2_rows"]
#         # ["limit_2_rows", "limit_3_rows"]: ["select_value"]
#     }
#     linear_dag = TransformerDAG(dag_nodes, linear_edges)
#
#     result_linears = build_and_execute_dag_pipeline(
#         linear_dag,
#         initial_data_frame_name="csv",
#         client=client,
#     )
#
#     for result_linear in result_linears:
#         for row in result_linear:
#             print(row.values)
#         print("=================")
