"""
TurboFrame Core - The main DataFrame class.
Wraps PyArrow Table with a fluent API and parallel execution.
"""

import time
import warnings
from typing import List, Dict, Optional, Union, Callable, Any

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq

try:
    from turboframe.parallel import ParallelEngine
except ImportError:
    from .parallel import ParallelEngine


class GroupByResult:
    """Holds groupby state until .agg() is called."""

    def __init__(self, frame: "TurboFrame", key_columns: List[str]):
        self._frame = frame
        self._key_columns = key_columns

    def agg(self, agg_specs: Dict[str, str]) -> "TurboFrame":
        """
        Execute parallel aggregation.

        Args:
            agg_specs: {"column": "sum|mean|min|max|count|std|var|median|nunique"}

        Returns:
            TurboFrame with aggregated results.

        Example:
            tf.groupby("region").agg({"sales": "sum", "quantity": "mean"})
        """
        engine = self._frame._engine
        result_table = engine.parallel_groupby(
            self._frame._table, self._key_columns, agg_specs
        )
        return TurboFrame._from_table(result_table, engine=engine)

    def count(self) -> "TurboFrame":
        """Shorthand: count rows per group."""
        first_non_key = [
            c for c in self._frame.columns if c not in self._key_columns
        ]
        col = first_non_key[0] if first_non_key else self._key_columns[0]
        return self.agg({col: "count"})

    def sum(self, columns: Optional[List[str]] = None) -> "TurboFrame":
        """Shorthand: sum specified columns per group."""
        cols = columns or [
            c for c in self._frame.columns if c not in self._key_columns
        ]
        return self.agg({c: "sum" for c in cols})


class TurboFrame:
    """
    A fast, parallel DataFrame built on PyArrow.

    Key design principles:
    - Zero-copy where possible (Arrow columnar format)
    - Parallel groupby via hash partitioning
    - Lazy-ish: filters push down, reads are optimized
    - Compatible: easy conversion to/from pandas, Spark, Polars
    """

    def __init__(
        self,
        data: Union[pa.Table, dict, "pd.DataFrame", None] = None,
        engine: Optional[ParallelEngine] = None,
        max_workers: Optional[int] = None,
    ):
        self._engine = engine or ParallelEngine(max_workers=max_workers)

        if data is None:
            self._table = pa.table({})
        elif isinstance(data, pa.Table):
            self._table = data
        elif isinstance(data, dict):
            self._table = pa.table(data)
        else:
            # Assume pandas DataFrame
            self._table = pa.Table.from_pandas(data)

    @classmethod
    def _from_table(
        cls, table: pa.Table, engine: Optional[ParallelEngine] = None
    ) -> "TurboFrame":
        """Internal: create TurboFrame from existing Arrow table."""
        tf = cls.__new__(cls)
        tf._table = table
        tf._engine = engine or ParallelEngine()
        return tf

    # ══════════════════════════════════════════════════════════════════
    # DATA READERS
    # ══════════════════════════════════════════════════════════════════

    @classmethod
    def read_delta(
        cls,
        path: str,
        columns: Optional[List[str]] = None,
        version: Optional[int] = None,
        storage_options: Optional[Dict[str, str]] = None,
        max_workers: Optional[int] = None,
    ) -> "TurboFrame":
        """
        Read a Delta table into TurboFrame. Uses deltalake (Rust-based) for speed.

        This is the FAST PATH:
        Delta (Parquet) → Arrow (zero-copy) → TurboFrame

        Args:
            path: Local path or cloud path (abfss://, s3://, gs://)
            columns: Only read these columns (column pruning = huge speedup)
            version: Read a specific Delta version (time travel)
            storage_options: Cloud credentials dict.
                For Fabric/ABFSS: usually auto-detected from notebook context.
            max_workers: Override CPU core count for parallelism.

        Example:
            # Fabric notebook
            tf = TurboFrame.read_delta(
                "abfss://container@account.dfs.core.windows.net/tables/sales"
            )

            # Local / Databricks
            tf = TurboFrame.read_delta("/mnt/delta/sales", columns=["region", "amount"])
        """
        try:
            from deltalake import DeltaTable
        except ImportError:
            raise ImportError(
                "Install deltalake: pip install deltalake\n"
                "On Fabric notebooks: %pip install deltalake"
            )

        t0 = time.perf_counter()

        # Build DeltaTable handle
        kwargs = {}
        if storage_options:
            kwargs["storage_options"] = storage_options
        if version is not None:
            kwargs["version"] = version

        dt = DeltaTable(path, **kwargs)

        # Read to Arrow — this is the fast path (Parquet → Arrow = zero-copy)
        arrow_table = dt.to_pyarrow_table(columns=columns)

        elapsed = time.perf_counter() - t0
        rows = len(arrow_table)
        cols = len(arrow_table.schema)
        print(f"[TurboFrame] Read {rows:,} rows × {cols} cols in {elapsed:.2f}s")

        return cls(data=arrow_table, max_workers=max_workers)

    @classmethod
    def read_parquet(
        cls,
        path: str,
        columns: Optional[List[str]] = None,
        max_workers: Optional[int] = None,
    ) -> "TurboFrame":
        """Read Parquet file(s) into TurboFrame."""
        t0 = time.perf_counter()
        table = pq.read_table(path, columns=columns)
        elapsed = time.perf_counter() - t0
        print(f"[TurboFrame] Read {len(table):,} rows in {elapsed:.2f}s from parquet")
        return cls(data=table, max_workers=max_workers)

    @classmethod
    def from_pandas(cls, df, max_workers: Optional[int] = None) -> "TurboFrame":
        """Convert pandas DataFrame to TurboFrame."""
        return cls(data=df, max_workers=max_workers)

    @classmethod
    def from_spark(cls, spark_df, max_workers: Optional[int] = None) -> "TurboFrame":
        """
        Convert Spark DataFrame to TurboFrame.
        Collects to driver — use for datasets that fit in driver memory.
        """
        pandas_df = spark_df.toPandas()
        return cls(data=pandas_df, max_workers=max_workers)

    @classmethod
    def from_dict(cls, data: dict, max_workers: Optional[int] = None) -> "TurboFrame":
        """Create TurboFrame from a dictionary of columns."""
        return cls(data=data, max_workers=max_workers)

    # ══════════════════════════════════════════════════════════════════
    # PROPERTIES
    # ══════════════════════════════════════════════════════════════════

    @property
    def shape(self) -> tuple:
        return (len(self._table), len(self._table.schema))

    @property
    def columns(self) -> List[str]:
        return self._table.column_names

    @property
    def dtypes(self) -> Dict[str, pa.DataType]:
        return {f.name: f.type for f in self._table.schema}

    @property
    def schema(self) -> pa.Schema:
        return self._table.schema

    @property
    def arrow_table(self) -> pa.Table:
        """Access the underlying PyArrow table directly."""
        return self._table

    def __len__(self) -> int:
        return len(self._table)

    def __repr__(self) -> str:
        rows, cols = self.shape
        return f"TurboFrame({rows:,} rows × {cols} cols, workers={self._engine.max_workers})"

    # ══════════════════════════════════════════════════════════════════
    # COLUMN OPERATIONS
    # ══════════════════════════════════════════════════════════════════

    def select(self, columns: List[str]) -> "TurboFrame":
        """Select specific columns (column pruning)."""
        table = self._table.select(columns)
        return TurboFrame._from_table(table, engine=self._engine)

    def drop(self, columns: Union[str, List[str]]) -> "TurboFrame":
        """Drop columns."""
        if isinstance(columns, str):
            columns = [columns]
        keep = [c for c in self.columns if c not in columns]
        return self.select(keep)

    def rename(self, mapping: Dict[str, str]) -> "TurboFrame":
        """Rename columns: {"old_name": "new_name"}."""
        new_names = [mapping.get(c, c) for c in self.columns]
        table = self._table.rename_columns(new_names)
        return TurboFrame._from_table(table, engine=self._engine)

    def with_column(self, name: str, values) -> "TurboFrame":
        """Add or replace a column."""
        if isinstance(values, (list, tuple)):
            arr = pa.array(values)
        elif isinstance(values, pa.Array):
            arr = values
        elif isinstance(values, pa.ChunkedArray):
            arr = values
        else:
            # Scalar — broadcast
            arr = pa.array([values] * len(self))

        if name in self.columns:
            idx = self.columns.index(name)
            table = self._table.set_column(idx, name, arr)
        else:
            table = self._table.append_column(name, arr)
        return TurboFrame._from_table(table, engine=self._engine)

    # ══════════════════════════════════════════════════════════════════
    # FILTERING
    # ══════════════════════════════════════════════════════════════════

    def filter(self, expression: str) -> "TurboFrame":
        """
        Filter rows using a simple expression string.

        Supported: >, <, >=, <=, ==, !=
        Examples:
            tf.filter("amount > 1000")
            tf.filter("region == 'East'")
            tf.filter("quantity != 0")
        """
        col_name, op, value = self._parse_filter_expr(expression)

        col_data = self._table.column(col_name)

        # Type-cast the value
        if pa.types.is_integer(col_data.type):
            val = int(value)
        elif pa.types.is_floating(col_data.type):
            val = float(value)
        elif pa.types.is_string(col_data.type) or pa.types.is_large_string(col_data.type):
            val = value.strip("'\"")
        else:
            val = value.strip("'\"")

        ops = {
            ">": pc.greater,
            "<": pc.less,
            ">=": pc.greater_equal,
            "<=": pc.less_equal,
            "==": pc.equal,
            "!=": pc.not_equal,
        }

        mask = ops[op](col_data, val)
        filtered = pc.filter(self._table, mask)
        return TurboFrame._from_table(filtered, engine=self._engine)

    def filter_arrow(self, mask: pa.Array) -> "TurboFrame":
        """Filter using a pre-computed boolean Arrow array."""
        filtered = pc.filter(self._table, mask)
        return TurboFrame._from_table(filtered, engine=self._engine)

    def filter_parallel(self, predicate_fn: Callable) -> "TurboFrame":
        """
        Filter with a custom function, executed in parallel across chunks.
        predicate_fn takes a pa.Table chunk and returns a boolean pa.Array.
        """
        result = self._engine.parallel_filter(self._table, predicate_fn)
        return TurboFrame._from_table(result, engine=self._engine)

    @staticmethod
    def _parse_filter_expr(expr: str):
        """Parse 'column op value' expressions."""
        for op in [">=", "<=", "!=", "==", ">", "<"]:
            if op in expr:
                parts = expr.split(op, 1)
                return parts[0].strip(), op, parts[1].strip()
        raise ValueError(f"Cannot parse filter expression: {expr}")

    # ══════════════════════════════════════════════════════════════════
    # SORTING
    # ══════════════════════════════════════════════════════════════════

    def sort(
        self, by: Union[str, List[str]], ascending: Union[bool, List[bool]] = True
    ) -> "TurboFrame":
        """Sort by one or more columns."""
        if isinstance(by, str):
            by = [by]
        if isinstance(ascending, bool):
            ascending = [ascending] * len(by)

        sort_keys = [
            (col, "ascending" if asc else "descending")
            for col, asc in zip(by, ascending)
        ]
        indices = pc.sort_indices(self._table, sort_keys=sort_keys)
        sorted_table = self._table.take(indices)
        return TurboFrame._from_table(sorted_table, engine=self._engine)

    # ══════════════════════════════════════════════════════════════════
    # GROUPBY
    # ══════════════════════════════════════════════════════════════════

    def groupby(self, keys: Union[str, List[str]]) -> GroupByResult:
        """
        Parallel GroupBy. Returns a GroupByResult, call .agg() to execute.

        How it works internally:
        1. Hash-partition the data across N CPU cores by key
        2. Each core aggregates its partition independently
        3. Results are concatenated (no merge needed — hash guarantees no overlap)

        Example:
            result = tf.groupby("region").agg({"sales": "sum", "qty": "mean"})
            result = tf.groupby(["region", "year"]).agg({"sales": "sum"})
        """
        if isinstance(keys, str):
            keys = [keys]
        return GroupByResult(self, keys)

    # ══════════════════════════════════════════════════════════════════
    # HEAD / TAIL / SAMPLE
    # ══════════════════════════════════════════════════════════════════

    def head(self, n: int = 5) -> "TurboFrame":
        return TurboFrame._from_table(self._table.slice(0, n), engine=self._engine)

    def tail(self, n: int = 5) -> "TurboFrame":
        start = max(0, len(self) - n)
        return TurboFrame._from_table(self._table.slice(start), engine=self._engine)

    def sample(self, n: int = 5, seed: Optional[int] = None) -> "TurboFrame":
        import numpy as np
        rng = np.random.default_rng(seed)
        indices = rng.choice(len(self), size=min(n, len(self)), replace=False)
        indices_sorted = pa.array(sorted(indices))
        return TurboFrame._from_table(
            self._table.take(indices_sorted), engine=self._engine
        )

    # ══════════════════════════════════════════════════════════════════
    # JOINS
    # ══════════════════════════════════════════════════════════════════

    def join(
        self,
        other: "TurboFrame",
        on: str,
        how: str = "inner",
    ) -> "TurboFrame":
        """
        Join two TurboFrames.
        how: "inner", "left", "right", "full outer"
        """
        # Map friendly names to PyArrow's join types
        join_map = {
            "inner": "inner",
            "left": "left outer",
            "right": "right outer",
            "outer": "full outer",
            "full": "full outer",
            "left outer": "left outer",
            "right outer": "right outer",
            "full outer": "full outer",
        }
        join_type = join_map.get(how, how)

        result = self._table.join(other._table, keys=on, join_type=join_type)
        return TurboFrame._from_table(result, engine=self._engine)

    # ══════════════════════════════════════════════════════════════════
    # COMPUTE / TRANSFORM
    # ══════════════════════════════════════════════════════════════════

    def apply_parallel(self, func: Callable) -> "TurboFrame":
        """
        Apply a transformation function in parallel across row chunks.
        func: takes a pa.Table → returns a pa.Table
        """
        result = self._engine.parallel_apply(self._table, func)
        return TurboFrame._from_table(result, engine=self._engine)

    # ══════════════════════════════════════════════════════════════════
    # CONVERSIONS
    # ══════════════════════════════════════════════════════════════════

    def to_pandas(self):
        """Convert to pandas DataFrame. Zero-copy where possible."""
        return self._table.to_pandas(self_destruct=False)

    def to_polars(self):
        """Convert to Polars DataFrame."""
        try:
            import polars as pl
            return pl.from_arrow(self._table)
        except ImportError:
            raise ImportError("Install polars: pip install polars")

    def to_spark(self, spark_session):
        """Convert to Spark DataFrame."""
        pdf = self.to_pandas()
        return spark_session.createDataFrame(pdf)

    def to_arrow(self) -> pa.Table:
        """Return the underlying PyArrow table."""
        return self._table

    def to_dict(self) -> Dict[str, list]:
        """Convert to Python dict of lists."""
        return {col: self._table.column(col).to_pylist() for col in self.columns}

    def to_parquet(self, path: str, compression: str = "snappy"):
        """Write to Parquet file."""
        pq.write_table(self._table, path, compression=compression)
        print(f"[TurboFrame] Wrote {len(self):,} rows to {path}")

    def to_delta(self, path: str, mode: str = "overwrite", storage_options: Optional[Dict] = None):
        """Write back to Delta table."""
        try:
            from deltalake import write_deltalake
        except ImportError:
            raise ImportError("Install deltalake: pip install deltalake")

        kwargs = {}
        if storage_options:
            kwargs["storage_options"] = storage_options

        write_deltalake(path, self._table, mode=mode, **kwargs)
        print(f"[TurboFrame] Wrote {len(self):,} rows to delta: {path}")

    # ══════════════════════════════════════════════════════════════════
    # DISPLAY
    # ══════════════════════════════════════════════════════════════════

    def show(self, n: int = 10):
        """Pretty-print first N rows."""
        sub = self._table.slice(0, min(n, len(self)))
        print(sub.to_pandas().to_string(index=False))

    def describe(self) -> Dict[str, Dict[str, Any]]:
        """Basic statistics for numeric columns."""
        stats = {}
        for col_name in self.columns:
            col = self._table.column(col_name)
            if pa.types.is_integer(col.type) or pa.types.is_floating(col.type):
                stats[col_name] = {
                    "count": pc.count(col).as_py(),
                    "mean": pc.mean(col).as_py(),
                    "min": pc.min(col).as_py(),
                    "max": pc.max(col).as_py(),
                    "std": pc.stddev(col).as_py(),
                    "null_count": col.null_count,
                }
        return stats

    # ══════════════════════════════════════════════════════════════════
    # UTILITY
    # ══════════════════════════════════════════════════════════════════

    def unique(self, column: str) -> list:
        """Get unique values of a column."""
        return pc.unique(self._table.column(column)).to_pylist()

    def value_counts(self, column: str) -> "TurboFrame":
        """Count occurrences of each value in a column."""
        result = pc.value_counts(self._table.column(column))
        # Convert StructArray to table
        values = [x["values"].as_py() for x in result]
        counts = [x["counts"].as_py() for x in result]
        table = pa.table({column: values, "count": counts})
        return TurboFrame._from_table(table, engine=self._engine)

    def is_empty(self) -> bool:
        return len(self) == 0

    def count_nulls(self) -> Dict[str, int]:
        """Count nulls per column."""
        return {col: self._table.column(col).null_count for col in self.columns}

    def drop_nulls(self, columns: Optional[List[str]] = None) -> "TurboFrame":
        """Drop rows with nulls in specified columns (or all)."""
        table = self._table.drop_null()
        return TurboFrame._from_table(table, engine=self._engine)

    def __del__(self):
        """Clean up parallel engine on garbage collection."""
        try:
            self._engine.shutdown()
        except Exception:
            pass
