from typing import Callable, Iterator, List, Union
from .operation import (
    Operation, SourceOp, MapOp, FilterOp, SelectOp, LimitOp,
    Row, FunctionWrapper, DataFrame
)
import inspect
# from .operation import FunctionWrapper

# ===== RemoteDataFrame 抽象接口 =====
class RemoteDataFrame:
    def get_schema(self) -> str:
        raise NotImplementedError

    def get_schema_uri(self) -> str:
        raise NotImplementedError

    def map(self, f: Callable[[Row], Row]) -> 'RemoteDataFrame':
        raise NotImplementedError

    def filter(self, f: Callable[[Row], bool]) -> 'RemoteDataFrame':
        raise NotImplementedError

    def select(self, *columns: str) -> 'RemoteDataFrame':
        raise NotImplementedError

    def limit(self, n: int) -> 'RemoteDataFrame':
        raise NotImplementedError

    def reduce(self, f: Callable[[Row, Row], Row]) -> 'RemoteDataFrame':
        raise NotImplementedError

    def foreach(self, f: Callable[[Row], None]) -> None:
        raise NotImplementedError

    def collect(self) -> List[Row]:
        raise NotImplementedError
    

# ===== RemoteDataFrame 实现类 =====
class RemoteDataFrameImpl(RemoteDataFrame):
    def __init__(self, data_frame_name: str, client, operation: Operation = None):
        self.data_frame_name = data_frame_name
        self.client = client
        self.operation = operation or SourceOp()

    def map(self, func) -> 'RemoteDataFrameImpl':
        wrapper = FunctionWrapper(func)
        return RemoteDataFrameImpl(self.data_frame_name, self.client, MapOp(wrapper, self.operation))

    def filter(self, func) -> 'RemoteDataFrameImpl':
        wrapper = FunctionWrapper(func)
        return RemoteDataFrameImpl(self.data_frame_name, self.client, FilterOp(wrapper, self.operation))

    def select(self, *columns: str) -> 'RemoteDataFrameImpl':
        return RemoteDataFrameImpl(self.data_frame_name, self.client, SelectOp(self.operation, *columns))

    def limit(self, n: int) -> 'RemoteDataFrameImpl':
        return RemoteDataFrameImpl(self.data_frame_name, self.client, LimitOp(n, self.operation))

    def reduce(self, func) -> 'RemoteDataFrameImpl':
        raise NotImplementedError("reduce is not yet supported")

    def foreach(self, func) -> None:
        for row in self._records():
            func(row)

    def collect(self) -> List[Row]:
        return list(self._records())

    def getSchema(self) -> str:
        return self.client.get_schema(self.data_frame_name)

    def getSchemaURI(self) -> str:
        return self.client.get_schema_uri(self.data_frame_name)

    def _records(self) -> Iterator[Row]:
        # 尝试根据客户端类型判断调用方式
        if hasattr(self.client, "get_rows"):
            try:
                result = self.client.get_rows(self.data_frame_name, self._get_operation_payload())
                for row_data in result:
                    yield Row(row_data)
            except Exception as e:
                raise RuntimeError(f"Failed to fetch rows from client: {e}")
        else:
            raise RuntimeError("Unsupported client type: missing 'get_rows' method")

    def _get_operation_payload(self) -> Union[str, dict]:
        """
        根据 client 类型返回合适的 operation 表达方式。
        """
        if self._is_remote_client():
            # print("It is remote")
            return self.operation.to_json()  # 字典（远程）
        else:
            # print("It is local")
            return self.operation.to_json_string()  # 字符串（本地）

    def _is_remote_client(self) -> bool:
        """
        判断是否为远程 ArrowFlightClient。
        """
        return hasattr(self.client, "token")

    def execute(self, op: Operation) -> List[Row]:
        """
        立即执行传入的 Operation，并返回结果 Row 列表。
        """
        temp_df = RemoteDataFrameImpl(self.data_frame_name, self.client, op)
        return temp_df.collect()
    # def collect_data(self) -> List[Row]:
    #     return self.collect()


# ===== 可选扩展：分组后聚合结果（占位） =====

class GroupedDataFrame:
    def __init__(self, remote_df: RemoteDataFrameImpl):
        self.remote_df = remote_df

    def max(self, column: str) -> RemoteDataFrameImpl:
        raise NotImplementedError("Aggregation not implemented")

    # 可扩展更多自定义聚合方法



