from typing import Any, Optional
from mloda.core.abstract_plugins.components.data_types import DataType
from mloda.provider import BaseMergeEngine
from mloda.provider import BaseFilterEngine, BaseMaskEngine
from mloda_plugins.compute_framework.base_implementations.pyarrow.pyarrow_merge_engine import PyArrowMergeEngine
from mloda_plugins.compute_framework.base_implementations.pyarrow.pyarrow_filter_engine import PyArrowFilterEngine
from mloda_plugins.compute_framework.base_implementations.pyarrow.pyarrow_mask_engine import (
    PyArrowMaskEngine,
)
import pyarrow as pa

from mloda.user import FeatureName
from mloda.provider import ComputeFramework


try:
    import pandas as pd
except ImportError:
    pd = None


class PyArrowTable(ComputeFramework):
    @staticmethod
    def is_available() -> bool:
        """Check if PyArrow is installed and available."""
        try:
            import pyarrow  # noqa: F401

            return True
        except ImportError:
            return False

    @classmethod
    def expected_data_framework(cls) -> Any:
        return pa.Table

    @classmethod
    def merge_engine(cls) -> type[BaseMergeEngine]:
        return PyArrowMergeEngine

    @classmethod
    def filter_engine(cls) -> type[BaseFilterEngine]:
        return PyArrowFilterEngine

    @classmethod
    def mask_engine(cls) -> type[BaseMaskEngine]:
        return PyArrowMaskEngine

    def select_data_by_column_names(
        self, data: Any, selected_feature_names: set[FeatureName], column_ordering: Optional[str] = None
    ) -> Any:
        column_names = set(data.schema.names)
        _selected_feature_names = self.identify_naming_convention(
            selected_feature_names, column_names, ordering=column_ordering
        )
        return data.select([f for f in _selected_feature_names])

    def _extract_column_names(self, data: Any) -> set[str]:
        return set(data.schema.names)

    def _extract_column_dtype(self, data: Any, column_name: str) -> str | None:
        if column_name in data.schema.names:
            return str(data.schema.field(column_name).type)
        return None

    def _extract_column_data_type(self, data: Any, column_name: str) -> Optional[DataType]:
        if column_name not in data.schema.names:
            return None
        return DataType.from_arrow_type_safe(data.schema.field(column_name).type)

    def transform(
        self,
        data: Any,
        feature_names: set[str],
    ) -> Any:
        transformed_data = self.apply_compute_framework_transformer(data)
        if transformed_data is not None:
            return transformed_data

        if isinstance(data, dict):
            """Initial data: Transform dict to table"""
            return pa.table(data)

        if isinstance(data, pa.ChunkedArray) or isinstance(data, pa.Array):
            """Added data: Add column to table"""
            if len(feature_names) == 1:
                return self.data.append_column(next(iter(feature_names)), data)
            raise ValueError(f"Only one feature can be added at a time: {feature_names}")

        raise ValueError(f"Data {type(data)} is not supported by {self.__class__.__name__}")
