import json
import uuid
import pyarrow as pa
import pyarrow.flight as flight
from typing import List, Dict, Any, Iterator, Union

from .auth_middle_ware import AuthMiddlewareFactory
from .blob import Blob
import re

from .code_utils import encode_with_map
from .dftpTicket import LoginTicket
from .operation import Operation, SourceOp, FilterOp, MapOp, LimitOp, SelectOp, TransformerNode
from .remote_dataframe import RemoteDataFrameImpl
from arrow_flight_project.udf.dag import TransformerDAG, Row, SourceNode, ConcreteUDFFunction, LimitNode, SelectNode
import ssl

# tls_cert_path = "/Users/guoliangzhu/PycharmProjects/pythonProject8/arrow_flight_project/faird"  # 修改为实际的 CA 根证书路径
# with open(tls_cert_path, "rb") as f:
#     root_cert = f.read()
#
# # 创建加密上下文（可选）
# tls_context = ssl.create_default_context(cadata=root_cert.decode())
from .simple_auth_handler import PairAuthHandler


class ArrowFlightClient:
    def __init__(self, host: str, port: int, use_tls: bool = False, ca_cert_path: str = None):
        # self.token = str(uuid.uuid4())
        # self.client = flight.FlightClient(flight.Location.for_grpc_tcp(host, port))
        # print(f"[INIT] 连接到 {host}:{port}\ntoken: {self.token}")
        self.auth_token = None
        self.token = str(uuid.uuid4())

        self.location = f"grpc+tcp://{host}:{port}" if not use_tls else f"grpc+tls://{host}:{port}"

        if use_tls:
            if not ca_cert_path:
                raise ValueError("使用 TLS 连接时必须提供 ca_cert_path")
            with open(ca_cert_path, "rb") as f:
                root_cert = f.read()
            self.client = flight.FlightClient(
                flight.Location.for_grpc_tls(host, port),
                tls_root_certs=root_cert
            )
            print(f"[INIT] 使用 TLS 连接到 {host}:{port}")
        else:
            self.client = flight.FlightClient(flight.Location.for_grpc_tcp(host, port))
            auth_handler = PairAuthHandler("Admin", "Admin")
            self.client.authenticate(auth_handler)

            print(f"[INIT] 使用非加密连接到 {host}:{port}")


    def login(self, credentials: Dict[str, str]):
        print(f"[LOGIN] 用户名: {credentials.get('username')}")
        payload = pa.RecordBatch.from_arrays(
            [pa.array([json.dumps(credentials).encode('utf-8')], type=pa.string())],
            schema=pa.schema([('credentials', pa.string())])
        )
        sink = pa.BufferOutputStream()
        with pa.ipc.new_stream(sink, payload.schema) as writer:
            writer.write_batch(payload)
        self._do_action(f"login.{self.token}", sink.getvalue().to_pybytes())
        print("[LOGIN] 成功")

    @classmethod
    def connect(cls, url: str, username: str, password: str, use_tls: bool = False, ca_cert_path: str = None) -> "ArrowFlightClient":
        """从 URL 和用户名密码构造并登录 client"""
        match = re.match(r"^(\w+://)?(?P<host>[^:]+):(?P<port>\d+)$", url)
        if not match:
            raise ValueError(f"无效 URL 格式: {url}")
        host = match.group("host")
        port = int(match.group("port"))
        client = ArrowFlightClient(host, port)
        # flight.Location.for_grpc_tcp(host, port)
        # client = flight.FlightClient(flight.Location.for_grpc_tcp(host, port))
        # auth_handler = PairAuthHandler("Admin", "Admin")
        # client.authenticate(auth_handler)

        # print("认证成功, token:", auth_handler.get_token())

        return client

    def open(self, name):
        return RemoteDataFrameImpl(name, self)

    # # 内部实现，通常不会被外部直接调用
    # def _do_action(self, action_type: str, body: bytes = b'') -> List[bytes]:
    #     """
    #     执行 Flight Action，并返回所有结果体的字节列表。
    #     """
    #     try:
    #         results = self.client.do_action(flight.Action(action_type, body))
    #         return [result.body.to_pybytes() for result in results]  # 确保转换为 Python bytes
    #     except flight.FlightInternalError as e:
    #         print(f"执行 Action '{action_type}' 失败，服务器返回内部错误: {e}")
    #         raise
    #     except Exception as e:
    #         print(f"执行 Action '{action_type}' 时发生未知错误: {e}")
    #         raise
    #
    #
    # def list_data_set_names(self) -> List[str]:
    #     """
    #     获取所有数据集的名称列表。
    #     """
    #     # print("获取数据集名称列表...")
    #     return self._parse_strings(self._do_action("listDataSetNames"))

    def _do_action(self, action_name: str, params: Dict[str, Any] = None, data: bytes = b'') -> List[bytes]:
        params = params or {}
        params['token'] = self.token  # 自动加 token
        body = encode_with_map(data, params)
        results = self.client.do_action(flight.Action(action_name, body))
        return [r.body.to_pybytes() for r in results]

    # from dftpTicket import DftpTicket, LoginTicket

    def list_data_set_names(self):
        """
        获取数据集列表，返回字符串列表
        """
        ticket = flight.Ticket(b"/listDataSets")
        reader = self.client.do_get(ticket)  # do_get 对应 Scala/Java get()

        data_sets = []

        while True:
            try:
                batch = reader.read_next_batch()
            except StopIteration:
                break
            # 遍历 RecordBatch
            for i in range(batch.num_rows):
                # 假设第一列是数据集名字
                ds_name = batch.column(0)[i].as_py()
                data_sets.append(ds_name)
        return data_sets

    # -------------------- List Dataset / DataFrame --------------------
    def get(self, url):
        """
        获取数据集列表，返回字符串列表
        """

        # json_str = f"""{"type": "SourceOp", "dataFrameName": "dftp://10.0.82.143:3101/"}"""

        data = {
            "type": "SourceOp",
            "dataFrameName": url
        }

        json_str = json.dumps(data)
        # json_str = f"""{"type": "SourceOp", "dataFrameName": ${url}}"""


        test_ticket = LoginTicket(json_str)

        print(test_ticket.type_id)
        print(test_ticket.ticket_content)
        encoded_ticket = test_ticket.encode_ticket()

        print(encoded_ticket)

        ticket = flight.Ticket(encoded_ticket)


        reader = self.client.do_get(ticket)  # do_get 对应 Scala/Java get()

        data_sets = []

        # while True:
        #     try:
        #         batch = reader.read_next_batch()
        #     except StopIteration:
        #         break
        #     # 遍历 RecordBatch
        #     for i in range(batch.num_rows):
        #         # 假设第一列是数据集名字
        #         ds_name = batch.column(0)[i].as_py()
        #         data_sets.append(ds_name)
        # return data_sets
        data_sets = []
        # reader = self.flight_client.do_get(self.ticket)

        for chunk in reader:
            batch = chunk.data  # 从 FlightStreamChunk 取出 RecordBatch
            data_sets.extend(batch.column(0).to_pylist())

        return data_sets

    def list_data_frame_names(self, ds_name: str) -> List[str]:
        """
        获取指定数据集下的所有 DataFrame 名称列表。
        """
        # print(f"获取数据集 '{ds_name}' 下的 DataFrame 名称列表...")
        return self._parse_strings(self._do_action(f"listDataFrameNames.{ds_name}"))

    def get_data_set_metadata(self, ds_name: str) -> str:
        """
        获取指定数据集的元数据（通常是 RDF/Jena 格式）。
        """
        # print(f"获取数据集 '{ds_name}' 的元数据...")
        return self._parse_single_string(self._do_action(f"getDataSetMetaData.{ds_name}"))

    def get_data_from_size(self, df_name: str) -> int:
        """
        获取指定 DataFrame 的行数。
        """
        # print(f"获取 DataFrame '{df_name}' 的行数...")
        return self._parse_single_long(self._do_action(f"getDataFrameSize.{df_name}"))

    def get_host_info(self) -> str:
        """
        获取服务器主机信息。
        """
        # print("获取主机信息...")
        return self._parse_single_string(self._do_action("getHostInfo"))

    def get_server_resource_info(self) -> str:
        """
        获取服务器资源使用情况。
        """
        # print("获取服务器资源信息...")
        return self._parse_single_string(self._do_action("getServerResourceInfo"))

    def get_rows(self, df_name: str, operation_dict: Dict[str, Any]) -> Iterator[Union[List[Any], Blob]]:
        """
        从 Flight 服务器获取数据行，支持数据操作和二进制列。

        Args:
            df_name (str): 要操作的 DataFrame 的名称。
                           注意: 如果你从 list_data_frame_names 获取，它可能是一个列表，需要取第一个元素。
            operation_dict (Dict[str, Any]): 表示数据操作的 Python 字典，会被序列化为 JSON 字符串。
                                             例如: {'op': 'select'}

        Yields:
            Union[List[Any], Blob]: 每行数据，如果最后一列是二进制类型，则包含 Blob 对象。
        """
        # request_id = str(uuid.uuid4())  # 生成唯一的请求ID
        # # print(f"开始获取行数据，请求ID: {request_id}")
        #
        # # 1. 准备并发送 do_put 请求 (上传操作描述符)
        # # 这个 schema 必须与服务器端 FlightProducerImpl.acceptPut 中期望的 schema 严格匹配
        # put_schema = pa.schema([
        #     ('dfName', pa.string()),
        #     ('userToken', pa.string()),
        #     ('DFOperation', pa.string())
        # ])
        #
        # # 确保 df_name 是一个字符串。如果它是一个列表，取第一个元素。
        # actual_df_name = df_name[0] if isinstance(df_name, list) and df_name else df_name
        # if not isinstance(actual_df_name, str):
        #     raise TypeError(f"df_name 必须是字符串或包含字符串的列表，但收到类型: {type(df_name)}")
        # if not actual_df_name:
        #     raise ValueError("df_name 不能为空字符串。")
        #
        # # 将 operation_dict 转换为 JSON 字符串
        # operation_json_string = json.dumps(operation_dict)
        #
        # # print(f"构建 put_batch: dfName='{actual_df_name}', token='{self.token}', operation='{operation_json_string}'")
        #
        # # 构建发送到服务器的 RecordBatch
        # # pa.array([single_value], type=pa.string()) 是将单个值放入一个单元素 PyArrow Array 的标准用法。
        # put_batch = pa.record_batch([
        #     pa.array([actual_df_name], type=pa.string()),
        #     pa.array([self.token], type=pa.string()),
        #     pa.array([operation_json_string], type=pa.string())
        # ], schema=put_schema)
        #
        #
        # # print("Flight descriptor path:", put_descriptor.path)
        #
        # # 使用 request_id 作为 FlightDescriptor 的路径
        # put_descriptor = flight.FlightDescriptor.for_path(request_id)
        #
        # # print(f"客户端发送 put_batch: \n{put_batch.to_pydict()}")
        # # print(f"客户端发送 put_descriptor 路径: {put_descriptor.path[0].decode()}") # 调试用
        #
        # # 执行 do_put 操作
        # writer, _ = self.client.do_put(put_descriptor, put_schema)
        # try:
        #     writer.write_batch(put_batch)
        #     # print("put_batch 写入完成。")
        # finally:
        #     writer.close()  # <-- 关键：确保关闭 writer，通知服务器数据已发送完毕
        #     # print("writer 已关闭，do_put 操作完成。")

        # 2. 调用 get_flight_info 获取数据流信息
        # 服务器的 getFlightInfo 会根据这个 request_id 从 requestMap 中查找信息
            # 确保 df_name 是字符串
        actual_df_name = df_name[0] if isinstance(df_name, list) and df_name else df_name

        action_payload = {
            'dfName': actual_df_name,
            'userToken': self.token,
            'type': operation_dict.get("type"),
            'input': operation_dict.get("input"),
            'args': operation_dict.get("args"),
        }
        body = json.dumps(action_payload).encode('utf-8')
        self._do_action(f"putRequest:{actual_df_name}:{self.token}", body)


        get_info_descriptor = flight.FlightDescriptor.for_path(self.token)

        # get_info_descriptor = flight.FlightDescriptor.for_path(request_id)
        # print(f"客户端请求 get_flight_info，描述符路径: {get_info_descriptor.path[0].decode()}")
        try:
            info = self.client.get_flight_info(get_info_descriptor)
            # print(f"成功获取到 FlightInfo，包含 {len(info.endpoints)} 个端点。")
        except flight.FlightInternalError as e:
            print(f"获取 FlightInfo 失败，服务器返回内部错误: {e}")
            print("这通常意味着服务器的 getFlightInfo 方法或其依赖的数据（例如 requestMap 中的条目）有问题。")
            raise

        # 3. 从 FlightInfo 中获取 Ticket 并执行 do_get
        if not info.endpoints:
            raise ValueError("FlightInfo 没有返回任何端点，无法获取数据。请检查服务器配置或数据可用性。")

        endpoint = info.endpoints[0]
        ticket = endpoint.ticket
        # print(f"获取到 Ticket: {ticket.ticket.decode()}，准备执行 do_get...")

        # do_get 用于从服务器下载数据流
        stream = self.client.do_get(ticket)

        # 4. 处理接收到的数据流
        response_schema = stream.schema
        # print(f"接收到的数据流 Schema:\n{response_schema}")

        # 检查最后一列是否是二进制类型，以便处理 Blob
        is_binary_last_column = False
        # if response_schema.num_fields > 0:
        last_field = response_schema.names[-1]
        is_binary_last_column = pa.types.is_binary(response_schema.field(last_field).type)

        num_rows_processed = 0
        try:
            # 遍历数据流中的所有批次
            for batch in stream:
                # 逐行处理批次中的数据
                for i in range(batch.data.num_rows):
                    values = [col[i].as_py() for col in batch.data.columns]
                    num_rows_processed += 1
                    if is_binary_last_column:
                        # 如果最后一列是二进制，将其包装成 Blob 对象
                        # 假设 Blob 的第一个参数是名称，第二个是数据迭代器
                        # 确保 values[0] 是有效的名称。如果只有一列或第一列不是名称，可能需要调整
                        name_for_blob = values[0] if values else f"blob_row_{num_rows_processed}"
                        blob = Blob(name_for_blob, iter([values[-1]]))  # values[-1] 是二进制数据
                        yield values[:-1] + [blob]  # 返回除最后一列外的所有值加上 Blob 对象
                    else:
                        # 如果不是二进制，直接返回所有列的值列表
                        yield values
            # print(f"数据流处理完成，共 {num_rows_processed} 行。")
        except Exception as e:
            print(f"处理数据流时发生错误: {e}")
            raise

    def close(self):
        """
        关闭 Flight 客户端连接。
        """
        # print("关闭 Flight 客户端连接...")
        self.client.close()
        # print("客户端连接已关闭。")

    # --- 辅助解析方法 ---
    def _parse_strings(self, result_bodies: List[bytes]) -> List[str]:
        """解析 action 返回的字符串列表。"""
        parsed_list = []
        for body in result_bodies:
            with pa.ipc.open_stream(body) as reader:
                for batch in reader:
                    # 假定字符串列表在第一列，且类型是 string
                    if batch.num_columns > 0 and pa.types.is_string(batch.schema.field(0).type):
                        parsed_list.extend(batch.columns[0].to_pylist())
                    else:
                        print(f"警告: _parse_strings 收到非预期 schema 或空批次: {batch.schema}")
        return parsed_list

    def _parse_single_string(self, result_bodies: List[bytes]) -> str:
        """解析 action 返回的单个字符串。"""
        if not result_bodies:
            raise ValueError("未收到任何结果体来解析单个字符串。")

        for body in result_bodies:
            with pa.ipc.open_stream(body) as reader:
                for batch in reader:
                    if batch.num_columns > 0 and batch.num_rows > 0 and pa.types.is_string(batch.schema.field(0).type):
                        return batch.columns[0][0].as_py()
        raise ValueError("无法从结果体中解析出单个字符串。")

    def _parse_single_long(self, result_bodies: List[bytes]) -> int:
        """解析 action 返回的单个长整型（Python int）。"""
        if not result_bodies:
            raise ValueError("未收到任何结果体来解析单个长整型。")

        for body in result_bodies:
            with pa.ipc.open_stream(body) as reader:
                for batch in reader:
                    if batch.num_columns > 0 and batch.num_rows > 0 and pa.types.is_int64(batch.schema.field(0).type):
                        return batch.columns[0][0].as_py()
        raise ValueError("无法从结果体中解析出单个长整型。")

    def execute_dag(
            self,
            dag: TransformerDAG,
    ) -> List[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()
                    # print(func_name)
                    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:
                        current_op = TransformerNode(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, self, 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
