"""
Parallel processing engine for TurboFrame.
Handles partitioning, multicore aggregation, and work distribution.
"""

import os
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from typing import List, Dict, Callable, Any, Optional, Tuple
from functools import partial

import pyarrow as pa
import pyarrow.compute as pc
import numpy as np


class ParallelEngine:
    """
    Manages parallel execution across CPU cores.
    Uses hash-based partitioning for groupby and thread pools for I/O.
    """

    def __init__(self, max_workers: Optional[int] = None):
        self.max_workers = max_workers or self._detect_optimal_workers()
        self._process_pool = None
        self._thread_pool = None

    def _detect_optimal_workers(self) -> int:
        cpu_count = os.cpu_count() or 4
        return max(1, cpu_count - 1)

    @property
    def process_pool(self) -> ProcessPoolExecutor:
        if self._process_pool is None:
            self._process_pool = ProcessPoolExecutor(max_workers=self.max_workers)
        return self._process_pool

    @property
    def thread_pool(self) -> ThreadPoolExecutor:
        if self._thread_pool is None:
            self._thread_pool = ThreadPoolExecutor(max_workers=self.max_workers)
        return self._thread_pool

    def shutdown(self):
        if self._process_pool:
            self._process_pool.shutdown(wait=False)
        if self._thread_pool:
            self._thread_pool.shutdown(wait=False)

    def hash_partition_table(
        self, table: pa.Table, key_column: str, n_partitions: Optional[int] = None
    ) -> List[pa.Table]:
        """
        Split a PyArrow table into N partitions by hashing the key column.
        All rows with same key land in same partition (critical for groupby).
        """
        n = n_partitions or self.max_workers
        key_array = table.column(key_column)

        if pa.types.is_string(key_array.type) or pa.types.is_large_string(key_array.type):
            hashes = np.array([hash(v.as_py()) % n for v in key_array], dtype=np.int64)
        elif pa.types.is_integer(key_array.type) or pa.types.is_floating(key_array.type):
            np_arr = key_array.to_numpy(zero_copy_only=False)
            hashes = np.mod(np_arr.astype(np.int64), n)
        else:
            hashes = np.array([hash(v.as_py()) % n for v in key_array], dtype=np.int64)

        partitions = []
        for pid in range(n):
            mask = hashes == pid
            if mask.any():
                indices = pa.array(np.where(mask)[0])
                partitions.append(table.take(indices))
        return partitions

    def parallel_groupby(
        self,
        table: pa.Table,
        key_columns: List[str],
        agg_specs: Dict[str, str],
    ) -> pa.Table:
        """
        Parallel groupby: hash-partition -> aggregate per partition -> merge.
        agg_specs: {"column_name": "sum|mean|min|max|count|..."}
        """
        if len(key_columns) == 1:
            partitions = self.hash_partition_table(table, key_columns[0])
        else:
            partitions = self.hash_partition_table(table, key_columns[0])

        if len(partitions) <= 1:
            return _single_groupby(table, key_columns, agg_specs)

        func = partial(_single_groupby, key_columns=key_columns, agg_specs=agg_specs)
        results = list(self.thread_pool.map(func, partitions))
        merged = pa.concat_tables([r for r in results if r is not None and len(r) > 0])
        return merged

    def parallel_filter(
        self, table: pa.Table, predicate_fn: Callable[[pa.Table], pa.Array]
    ) -> pa.Table:
        n = self.max_workers
        total_rows = len(table)
        chunk_size = max(1, total_rows // n)

        chunks = []
        for i in range(0, total_rows, chunk_size):
            end = min(i + chunk_size, total_rows)
            chunks.append(table.slice(i, end - i))

        def _apply_filter(chunk):
            mask = predicate_fn(chunk)
            return pc.filter(chunk, mask)

        results = list(self.thread_pool.map(_apply_filter, chunks))
        non_empty = [r for r in results if len(r) > 0]
        return pa.concat_tables(non_empty) if non_empty else table.slice(0, 0)

    def parallel_apply(
        self, table: pa.Table, func: Callable[[pa.Table], pa.Table]
    ) -> pa.Table:
        n = self.max_workers
        total_rows = len(table)
        chunk_size = max(1, total_rows // n)

        chunks = []
        for i in range(0, total_rows, chunk_size):
            end = min(i + chunk_size, total_rows)
            chunks.append(table.slice(i, end - i))

        results = list(self.thread_pool.map(func, chunks))
        return pa.concat_tables(results)


def _single_groupby(
    table: pa.Table,
    key_columns: List[str],
    agg_specs: Dict[str, str],
) -> pa.Table:
    """Perform groupby on a single partition using PyArrow native group_by."""
    if table is None or len(table) == 0:
        return table

    gb = table.group_by(key_columns)

    agg_map = {
        "sum": "sum",
        "mean": "mean",
        "min": "min",
        "max": "max",
        "count": "count",
        "std": "stddev",
        "var": "variance",
        "median": "approximate_median",
        "nunique": "count_distinct",
        "any": "any",
        "all": "all",
    }

    for col, agg_type in agg_specs.items():
        arrow_agg = agg_map.get(agg_type)
        if arrow_agg is None:
            raise ValueError(
                f"Unknown aggregation '{agg_type}'. Supported: {list(agg_map.keys())}"
            )
        gb = gb.aggregate([(col, arrow_agg)])

    return gb
