# from arrow_flight_project.udf.dag import SourceNode, LimitNode, SelectNode, ConcreteUDFFunction, TransformerDAG, \
#     build_and_execute_dag_pipeline
from arrow_flight_project.udf.dag import SourceNode, LimitNode, SelectNode, ConcreteUDFFunction, TransformerDAG
from client.arrow_flight_client import ArrowFlightClient
from client.remote_dataframe import RemoteDataFrameImpl
from client.operation import SourceOp, FilterOp, FunctionWrapper, MapOp, SelectOp, LimitOp
from udf.filters import filter_id_less_than_10
from udf.map import map_value_squre_10
from udf.custom_func import custom_func

import pyarrow as pa
import pyarrow.flight as fl


def test_dag(client):
    # 模拟节点图结构如下：
    source_node_1 = SourceNode("/csv/data_1.csv")
    source_node_2 = SourceNode('/csv/data_2.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)
    custom_func_wrapper = FunctionWrapper(custom_func)

    # 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,
        "custom_op": ConcreteUDFFunction(custom_func_wrapper),
        "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"],
    #     "source_2": ["map_op"],
    #     "map_op": ["filter_op"],
    #     "filter_op": ["limit_2_rows", "limit_3_rows"],
    # }
    linear_edges = {
        "source_1": ["custom_op"],
        "custom_op": ["map_op"],
        "map_op": ["limit_2_rows", "limit_3_rows"],
        # "source_2": ["map_op"],
        # "map_op": ["filter_op"],
        # "filter_op": ["limit_2_rows", "limit_3_rows"],

        # "filter_op": ["custom_op"],
        # "custom_op": ["limit_2_rows", "limit_3_rows"],
    }

    dag = TransformerDAG(dag_nodes, linear_edges)
    results = client.execute_dag(dag)

    for result in results:
        # print("=================")
        for row in result:
            print(row.values)
        print("\n")


def test_upload(client):

    # 上传文件-pandas-上传
    # client结构，用户实现xx接口，schema
    # 1. 定义 schema
    schema = pa.schema([
        ("name", pa.string()),
        ("value", pa.int64())
    ])

    # 2. 获取 writer
    descriptor = fl.FlightDescriptor.for_path("big_dataset")  # path[0] = dfName
    writer, metadata_reader = client.do_put(descriptor, schema)

    # 3. 生成批次数据
    batch = pa.record_batch([
        pa.array(["a", "b", "c", "d"]),
        pa.array([1, 2, 3, 4])
    ], schema=schema)

    writer.write_batch(batch)

    # 4. 结束上传
    writer.done_writing()

    # 5. 等待服务器ACK
    while True:
        msg = metadata_reader.read()
        if msg is None:
            break
        print("Server ACK:", bytes(msg).decode("utf-8"))


def main():
    try:
        # connect
        # client = None
        # client = ArrowFlightClient.connect("localhost:3101", "admin@instdb.cn", "admin001")
        client = ArrowFlightClient.connect("10.0.82.143:3105", "Admin", "Admin")
        #/Users/guoliangzhu/Downloads/Faird-main 3/src/main/resources/server.crt
        # client = ArrowFlightClient.connect("localhost:3101", "admin@instdb.cn", "admin001", use_tls=True, ca_cert_path="/Users/guoliangzhu/PycharmProjects/pythonProject8/arrow_flight_project/faird")

        # client = ArrowFlightClient.connect("dacp://10.0.82.143:3101", "admin", "admin")

        dataset_name = client.get("dftp://10.0.82.143:3105/")
        print("📂 数据集名称列表 (dataset_name):", dataset_name)

        dataframe_name = client.get("dftp://10.0.82.143:3105/"+dataset_name[0])
        print("📂 第一个数据集中所有的数据帧名称 (data_frame_name):", dataframe_name)

        dataframe_data = client.get("dftp://10.0.82.143:3105/"+dataset_name[0]+"/"+dataframe_name[0])
        print("📂 数据帧详细信息 (data_frame_data):", dataframe_data)


        # # "dacp://10.0.82.143:3101", UsernamePassword("admin", "admin")
        # # 获取数据集名称列表
        # dataset_name = client.list_data_set_names()
        # print("📂 数据集名称列表 (dataset_name):", dataset_name)
        #
        # # 获取第一个数据集的元信息
        # data_set_meta_data = client.get_data_set_metadata(dataset_name[0])
        # print("📄 数据集元信息 (data_set_meta_data):", data_set_meta_data)
        #
        # # 获取第一个数据集中所有数据帧名称
        # data_frame_name = client.list_data_frame_names(dataset_name[1])
        # print("📊 第一个数据集中所有的数据帧名称 (data_frame_name):", data_frame_name)
        #
        # # 获取第一个数据集的大小信息
        # dataset_size = client.get_data_from_size(dataset_name[0])
        # print("📦 数据集大小 (dataset_size):", dataset_size)
        #
        # # 获取客户端当前连接的主机信息
        # host_info = client.get_host_info()
        # print("🖥️ 主机信息 (host_info):", host_info)
        #
        # # 获取服务器资源（如CPU、内存等）信息
        # resource_info = client.get_server_resource_info()
        # print("📈 服务器资源信息 (resource_info):", resource_info)
        #
        # 测试非结构化数据帧
        # print("\n📂 非结构化数据帧 (/bin)：")
        # df_bin = client.open("/bin")
        # df_bin.foreach(lambda x: print(x.values))
        # #
        # # 测试结构化数据帧
        # print("\n📊 结构化数据帧 (csv) ，过滤 + 映射 + 限制 + 选择:")
        # df = client.open("/csv/data_1.csv")
        # # df.map(map_value_squre_10).limit(20).foreach(lambda x: print(x.values))
        # # result = df.filter(filter_id_less_than_10).map(map_value_squre_10).limit(5).select("value").collect()
        #
        # df.filter(filter_id_less_than_10). \
        #     map(map_value_squre_10). \
        #     limit(5). \
        #     select("_2"). \
        #     foreach(lambda x: print(x.values))
        #
        # # # result = df.filter(lambda x: x.get(0) < 10).map(lambda x: {"id": x.get(0), "value": x.get(1) * 10}).limit(5).select("value").collect()
        # #
        # # # 测试DAG执行
        # print("\n🧬 DAG 执行结果:")
        # test_dag(client)
        # test_upload(client.client)

    finally:
        if client:
            client.close()


if __name__ == "__main__":
    main()
