from abc import ABC, abstractmethod
from typing import Any, Dict, Iterator, List
import json
import inspect
import textwrap
import re
import ast
# from RemoteDataFrame import RemoteDataFrameImpl


class Row:
    def __init__(self, values: List[Any]):
        self.values = values

    def get(self, index: int) -> Any:
        return self.values[index]

    @staticmethod
    def from_seq(seq: List[Any]) -> 'Row':
        return Row(seq)


class Schema:
    def __init__(self, fields: List[str]):
        self.fields = fields

    def index_of(self, name: str):
        try:
            return self.fields.index(name)
        except ValueError:
            return None

    def select(self, *cols: str) -> 'Schema':
        return Schema([c for c in cols if c in self.fields])


class DataFrame:
    def __init__(self, schema: Schema, stream: Iterator[Row]):
        self.schema = schema
        self.stream = stream

def get_dataframe_by_stream(stream: Iterator[Row], schema: Schema) -> DataFrame:
    return DataFrame(schema, stream)


class FunctionWrapper:
    def __init__(self, impl):
        self.impl = impl

    def apply_to_input(self, value, interp=None):
        return self.impl(value)

    def to_json(self) -> Dict[str, Any]:
        try:
            # 获取函数源代码
            func_src = inspect.getsource(self.impl)

            # print("=========")
            # print(func_src)

            func_src = textwrap.dedent(func_src)

            # 用 AST 分析函数中用到的模块名（如 math, json 等）
            used_names = set()
            tree = ast.parse(func_src)
            func_name = None  #  提取函数名
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef) and func_name is None:
                    func_name = node.name
                elif isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
                    used_names.add(node.value.id)
                elif isinstance(node, ast.Name):
                    used_names.add(node.id)

            # 从函数所在模块中提取 import 行
            import_lines = []
            module = inspect.getmodule(self.impl)
            if module and hasattr(module, '__file__'):
                try:
                    with open(module.__file__, 'r', encoding='utf-8') as f:
                        module_src = f.read()
                    all_imports = re.findall(r'^\s*(import .+|from .+ import .+)', module_src, re.MULTILINE)
                    for line in all_imports:
                        for name in used_names:
                            if re.search(rf'\b{name}\b', line):
                                import_lines.append(line)
                                break  # 一行只加一次
                except Exception:
                    pass

            # 拼接调用语句
            if func_name:
                call_line = f"output_data = {func_name}()"
            else:
                call_line = "# [ERROR] Could not determine function name"

            full_code = '\n'.join(import_lines) + '\n\n' + func_src + '\n\n' + call_line

        except Exception as e:
            full_code = f"# Failed to extract source\n# Error: {e}"
        return {"type": "PYTHON_CODE", "code": full_code}
        # return {"type": "PYTHON_BIN", "functionId": "123", "functionName": "normalize", "whlPath":"/Users/guoliangzhu/PycharmProjects/pythonProject8/arrow_flight_project/link-0.1-py3-none-any.whl"}

    def whl_to_json(self) -> Dict[str, Any]:
        return {"type": "PYTHON_BIN", "functionID": "123", "functionName": "normalize", "whlPath":"/Users/guoliangzhu/Documents/faird-java-dir/faird-java/faird-core/target/test-classes/lib/link-0.1-py3-none-any.whl"}



    @staticmethod
    def from_json(obj: Dict[str, Any]) -> 'FunctionWrapper':
        code_str = obj["code"]
        local_ns = {}
        exec(code_str, {}, local_ns)
        func = next(v for v in local_ns.values() if callable(v))
        return FunctionWrapper(func)



# ==== Operation 基类与实现 ====

class Operation(ABC):
    @property
    @abstractmethod
    def operation_type(self) -> str:
        pass

    @abstractmethod
    def to_json(self) -> Dict[str, Any]:
        pass

    def to_json_string(self) -> str:
        return json.dumps(self.to_json())

    @abstractmethod
    def execute(self, df: DataFrame) -> DataFrame:
        pass

    @staticmethod
    def from_json_string(json_str: str) -> 'Operation':
        parsed = json.loads(json_str)
        op_type = parsed["type"]

        if op_type == "SourceOp":
            return SourceOp()

        # 对于有输入的操作，先递归解析其输入
        input_op = None
        if "input" in parsed:
            input_op = Operation.from_json_string(json.dumps(parsed["input"]))
        if op_type == "Map":
            # 正确地反序列化 function_wrapper
            func_wrapper_json = parsed["function"]
            func_wrapper = FunctionWrapper.from_json(func_wrapper_json)
            return MapOp(func_wrapper, input_op)
        elif op_type == "Filter":
            # 正确地反序列化 function_wrapper
            func_wrapper_json = parsed["function"]
            func_wrapper = FunctionWrapper.from_json(func_wrapper_json)
            return FilterOp(func_wrapper, input_op)
        elif op_type == "Limit":
            return LimitOp(parsed["args"][0], input_op)
        elif op_type == "Select":
            return SelectOp(input_op, *parsed["args"])
        elif op_type == "TransformerNode":
            # 正确地反序列化 function_wrapper
            func_wrapper_json = parsed["function"]
            func_wrapper = FunctionWrapper.from_json(func_wrapper_json)
            return TransformerNode(func_wrapper, input_op)
        else:
            raise ValueError(f"Unknown operation type: {op_type}")

class SourceOp(Operation):
    @property
    def operation_type(self) -> str:
        return "SourceOp"

    def to_json(self) -> Dict[str, Any]:
        return {"type": self.operation_type}

    def execute(self, df: DataFrame) -> DataFrame:
        return df


class MapOp(Operation):
    def __init__(self, function_wrapper: FunctionWrapper, input_op: Operation):
        self.function_wrapper = function_wrapper
        self.input = input_op

    @property
    def operation_type(self) -> str:
        return "Map"

    def to_json(self) -> Dict[str, Any]:
        return {
            "type": self.operation_type,
            "function": self.function_wrapper.to_json(),
            "input": self.input.to_json()
        }

    def execute(self, df: DataFrame) -> DataFrame:
        in_df = self.input.execute(df)
        new_stream = (self.function_wrapper.apply_to_input(row) for row in in_df.stream)
        return get_dataframe_by_stream(new_stream, in_df.schema)


class FilterOp(Operation):
    def __init__(self, function_wrapper: FunctionWrapper, input_op: Operation):
        self.function_wrapper = function_wrapper
        self.input = input_op

    @property
    def operation_type(self) -> str:
        return "Filter"

    def to_json(self) -> Dict[str, Any]:
        return {
            "type": self.operation_type,
            "function": self.function_wrapper.to_json(),
            "input": self.input.to_json()
        }

    def execute(self, df: DataFrame) -> DataFrame:
        in_df = self.input.execute(df)
        new_stream = (row for row in in_df.stream if self.function_wrapper.apply_to_input(row))
        return get_dataframe_by_stream(new_stream, in_df.schema)


class LimitOp(Operation):
    def __init__(self, limit: int, input_op: Operation):
        self.limit = limit
        self.input = input_op

    @property
    def operation_type(self) -> str:
        return "Limit"

    def to_json(self) -> Dict[str, Any]:
        return {
            "type": self.operation_type,
            "args": [self.limit],
            "input": self.input.to_json()
        }

    def execute(self, df: DataFrame) -> DataFrame:
        in_df = self.input.execute(df)
        new_stream = (row for idx, row in enumerate(in_df.stream) if idx < self.limit)
        return get_dataframe_by_stream(new_stream, in_df.schema)


class SelectOp(Operation):
    def __init__(self, input_op: Operation, *columns: str):
        self.input = input_op
        self.columns = list(columns)

    @property
    def operation_type(self) -> str:
        return "Select"

    def to_json(self) -> Dict[str, Any]:
        return {
            "type": self.operation_type,
            "args": self.columns,
            "input": self.input.to_json()
        }

    def execute(self, df: DataFrame) -> DataFrame:
        in_df = self.input.execute(df)
        new_schema = in_df.schema.select(*self.columns)
        new_stream = (
            Row.from_seq([
                row.get(in_df.schema.index_of(col))
                for col in self.columns
            ])
            for row in in_df.stream
        )
        return get_dataframe_by_stream(new_stream, new_schema)


class TransformerNode(Operation):
    def __init__(self, function_wrapper: FunctionWrapper, input_op: Operation):
        self.function_wrapper = function_wrapper
        self.input = input_op

    @property
    def operation_type(self) -> str:
        return "TransformerNode"

    def to_json(self) -> Dict[str, Any]:
        return {
            "type": self.operation_type,
            # "function": self.function_wrapper.to_json(),
            "function": self.function_wrapper.whl_to_json(),
            "input": self.input.to_json()
        }

    def execute(self, df: DataFrame) -> DataFrame:
        in_df = self.input.execute(df)
        transformed_stream = self.function_wrapper.apply_to_input(in_df.stream)
        return get_dataframe_by_stream(transformed_stream, in_df.schema)
