"""
Fast reads: column pruning, optional integer-column range filters, PyArrow predicate
pushdown (with consistent post-filters when range cannot be pushed), unified multi-file
scans with threaded Arrow decode and optional row-group pre-buffering, parallel per-file
reads, and tuned Arrow→pandas conversion.
"""

from __future__ import annotations

import os
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, List, Optional, Sequence, Tuple, Union

import pandas as pd
import pyarrow as pa

from snapparquet.dataset_expr import integer_range_expression, merge_dataset_expressions
from snapparquet.filters import FilterSpec, apply_filters_pandas, build_filter
from snapparquet.schema import intersect_columns, norm_col

PathLike = Union[str, Path]


def _parquet_read_format():
    """
    Parquet file format tuned for sequential scans: pre-buffer row groups where supported.

    Falls back to default format if options are unavailable (older PyArrow builds).
    """
    import pyarrow.dataset as ds

    try:
        scan = ds.ParquetFragmentScanOptions(pre_buffer=True)
        return ds.ParquetFileFormat(default_fragment_scan_options=scan)
    except Exception:
        return ds.ParquetFileFormat()


def _dataset_to_table(
    dataset: Any,
    *,
    columns: Optional[List[str]],
    filter_: Any,
) -> "pa.Table":
    """
    Materialize a dataset with multi-threaded decode and optional predicate pushdown.

    *filter_* may be ``None`` (read all rows, still respects *columns* projection).
    """
    kwargs: dict = {"columns": columns, "filter": filter_, "use_threads": True}
    try:
        return dataset.to_table(**kwargs)
    except TypeError:
        kwargs.pop("use_threads", None)
        return dataset.to_table(columns=columns, filter=filter_)


def _apply_unpushed_legacy_range(
    df: pd.DataFrame,
    *,
    range_ok: bool,
    expr_range: Any,
    fc: Optional[str],
    lo: Optional[int],
    hi: Optional[int],
) -> pd.DataFrame:
    """
    When tuple filters (or other expressions) use the Arrow dataset path but the legacy
    integer *filter_column* range was not pushed (non-integer column, pushdown disabled,
    etc.), apply the inclusive range in pandas so semantics match the pandas fallback path.
    """
    if (
        range_ok
        and fc is not None
        and lo is not None
        and hi is not None
        and expr_range is None
    ):
        return _filter_dataframe_range(df, fc, lo, hi)
    return df


def _supports_pandas_pyarrow_backend() -> bool:
    """``types_mapper=pd.ArrowDtype`` requires pandas 2.0+."""
    try:
        major = int(pd.__version__.split(".", 1)[0])
        return major >= 2
    except Exception:
        return False


def _chunked_dictionary_to_categorical(col: pa.ChunkedArray) -> Optional[pd.Categorical]:
    """Decode Arrow dictionary encoding to pandas ``Categorical``."""
    merged = col.combine_chunks()
    if not pa.types.is_dictionary(merged.type):
        return None
    try:
        indices = merged.indices.to_numpy(zero_copy_only=True)
    except Exception:
        indices = merged.indices.to_numpy(zero_copy_only=False)
    dictionary = merged.dictionary
    try:
        cats = dictionary.to_pylist()
    except Exception:
        cats = [dictionary[j].as_py() for j in range(len(dictionary))]
    return pd.Categorical.from_codes(indices, categories=cats)


def _dictionary_columns_to_categorical(df: pd.DataFrame, table: pa.Table) -> pd.DataFrame:
    """Replace dictionary-encoded columns with pandas categoricals (numpy backend path)."""
    if table is None or table.num_columns == 0:
        return df
    out = df
    mutated = False
    for i, field in enumerate(table.schema):
        if not pa.types.is_dictionary(field.type):
            continue
        name = norm_col(field.name)
        if name not in out.columns:
            continue
        try:
            cat = _chunked_dictionary_to_categorical(table.column(i))
            if cat is not None:
                if not mutated:
                    out = out.copy()
                    mutated = True
                out[name] = cat
        except Exception:
            continue
    return out


def _arrow_table_to_pandas(
    table: pa.Table,
    *,
    backend: str = "numpy",
    dictionary_as_categorical: bool = True,
) -> pd.DataFrame:
    """
    Convert Arrow table to pandas.

    * ``backend="numpy"`` (default): classic numpy-backed dtypes; optional dictionary → categorical.
    * ``backend="pyarrow"``: ``types_mapper=pd.ArrowDtype`` when pandas >= 2; else warns and falls back.
    """
    backend = (backend or "numpy").lower()
    if backend not in ("numpy", "pyarrow"):
        raise ValueError("backend must be 'numpy' or 'pyarrow', got {!r}".format(backend))

    if backend == "pyarrow" and not _supports_pandas_pyarrow_backend():
        warnings.warn(
            "snapparquet: backend='pyarrow' requires pandas>=2.0; using numpy-backed pandas.",
            UserWarning,
            stacklevel=2,
        )
        backend = "numpy"

    # self_destruct=True frees Arrow buffers during to_pandas; that invalidates *table*
    # before dictionary→categorical post-processing can read chunked arrays.
    allow_destruct = not (backend == "numpy" and dictionary_as_categorical)

    def _to_df() -> pd.DataFrame:
        if backend == "pyarrow":
            try:
                kw = dict(split_blocks=True, use_threads=True)
                if allow_destruct:
                    kw["self_destruct"] = True
                return table.to_pandas(types_mapper=pd.ArrowDtype, **kw)
            except TypeError:
                try:
                    return table.to_pandas(
                        types_mapper=pd.ArrowDtype,
                        split_blocks=True,
                        use_threads=True,
                    )
                except TypeError:
                    return table.to_pandas(types_mapper=pd.ArrowDtype, use_threads=True)
        try:
            kw = dict(split_blocks=True, use_threads=True)
            if allow_destruct:
                kw["self_destruct"] = True
            return table.to_pandas(**kw)
        except TypeError:
            try:
                return table.to_pandas(split_blocks=True, use_threads=True)
            except TypeError:
                try:
                    return table.to_pandas(use_threads=True)
                except TypeError:
                    return table.to_pandas()

    df = _to_df()
    if backend == "numpy" and dictionary_as_categorical:
        df = _dictionary_columns_to_categorical(df, table)
    return df


def _filter_dataframe_range(
    df: pd.DataFrame,
    column: str,
    low: int,
    high: int,
) -> pd.DataFrame:
    if df is None or len(df) == 0:
        return df
    col = norm_col(column)
    if col not in df.columns:
        return df
    v = pd.to_numeric(df[col], errors="coerce")
    return df[(v >= low) & (v <= high)].copy()


def _ensure_filter_column_in_list(
    columns: Optional[List[str]], filter_column: str
) -> Optional[List[str]]:
    if columns is None:
        return None
    fc = norm_col(filter_column)
    if fc in columns:
        return columns
    out = list(columns)
    out.append(fc)
    return out


def _ensure_filter_columns_in_list(
    columns: Optional[List[str]], filters: Optional[FilterSpec]
) -> Optional[List[str]]:
    """Append any filter tuple columns so dataset reads include predicate fields."""
    if columns is None or not filters:
        return columns
    have = {norm_col(c) for c in columns}
    out = list(columns)
    for col, _, _ in filters:
        n = norm_col(str(col))
        if n not in have:
            out.append(n)
            have.add(n)
    return out


def _paths_are_all_parquet(paths: Sequence[str]) -> bool:
    for p in paths:
        low = str(p).strip().lower().rstrip("/")
        if not low.endswith(".parquet"):
            return False
    return bool(paths)


def _try_read_parquet_unified(
    paths: Sequence[str],
    *,
    columns: Optional[Sequence[str]],
    filter_column: Optional[str],
    range_start: Optional[int],
    range_end: Optional[int],
    use_pushdown: bool,
    filters: Optional[FilterSpec] = None,
    backend: str = "numpy",
    dictionary_as_categorical: bool = True,
) -> Optional[pd.DataFrame]:
    """
    Single pyarrow.dataset scan over many Parquet files (or directories).
    Returns None on failure so callers can fall back to per-file reads.
    """
    plist = [str(p).strip() for p in paths if str(p).strip()]
    if not plist or not _paths_are_all_parquet(plist):
        return None
    try:
        import pyarrow.dataset as ds
    except Exception:
        return None

    cols = list(columns) if columns is not None else None
    lo = int(range_start) if range_start is not None else None
    hi = int(range_end) if range_end is not None else None
    range_ok = lo is not None and hi is not None and filter_column is not None
    fc = norm_col(filter_column) if filter_column else None

    if range_ok and fc:
        cols = _ensure_filter_column_in_list(cols, fc)
    cols = _ensure_filter_columns_in_list(cols, filters)

    expr_range = integer_range_expression(
        plist[0],
        cols,
        use_pushdown=use_pushdown,
        filter_column=filter_column,
        range_start=range_start,
        range_end=range_end,
    )
    expr_filters = build_filter(filters) if filters else None
    flt = merge_dataset_expressions(expr_range, expr_filters)

    try:
        dataset = ds.dataset(plist, format=_parquet_read_format())
        table = _dataset_to_table(dataset, columns=cols, filter_=flt)
        df = _arrow_table_to_pandas(
            table,
            backend=backend,
            dictionary_as_categorical=dictionary_as_categorical,
        )
        return _apply_unpushed_legacy_range(
            df,
            range_ok=range_ok,
            expr_range=expr_range,
            fc=fc,
            lo=lo,
            hi=hi,
        )
    except Exception:
        return None


def _default_io_workers(num_paths: int) -> int:
    if num_paths <= 2:
        return 1
    cpu = os.cpu_count() or 4
    # Per-file reads are I/O bound; modest oversubscription hides latency (local disk / object stores).
    return min(32, num_paths, max(8, cpu * 4))


def _finalize_parquet_dataframe(
    df: pd.DataFrame,
    *,
    range_ok: bool,
    fc: Optional[str],
    lo: Optional[int],
    hi: Optional[int],
    filters: Optional[FilterSpec],
) -> pd.DataFrame:
    """Apply legacy range + tuple filters in pandas (fallback path)."""
    if range_ok and fc is not None and lo is not None and hi is not None:
        df = _filter_dataframe_range(df, fc, lo, hi)
    if filters:
        df = apply_filters_pandas(df, filters)
    return df


def read_parquet(
    path: PathLike,
    *,
    columns: Optional[Sequence[str]] = None,
    filter_column: Optional[str] = None,
    range_start: Optional[int] = None,
    range_end: Optional[int] = None,
    use_pushdown: bool = True,
    filters: Optional[FilterSpec] = None,
    backend: str = "numpy",
    dictionary_as_categorical: bool = True,
) -> pd.DataFrame:
    """
    Read a Parquet file into a pandas DataFrame.

    * **Column pruning** — pass ``columns`` to decode only those fields.
    * **Legacy range** — ``filter_column`` + ``range_start`` / ``range_end`` (inclusive).
      When the column is integer-typed in the file, combined with ``filters`` into one
      ``pyarrow.dataset`` predicate when possible; otherwise applied in pandas.
    * **Tuple filters** — ``filters=[("col", "op", value), ...]`` for pushdown via
      :func:`snapparquet.filters.build_filter` (see module docstring).
    * **backend** — ``"numpy"`` (default) or ``"pyarrow"`` (pandas 2+ ``ArrowDtype``;
      falls back with a warning on older pandas).
    * **dictionary_as_categorical** — when ``backend="numpy"``, map Arrow dictionary
      columns to ``pandas.Categorical`` to reduce memory.
    """
    be = (backend or "numpy").strip().lower()
    if be not in ("numpy", "pyarrow"):
        raise ValueError("backend must be 'numpy' or 'pyarrow', got {!r}".format(backend))

    p = str(path).strip()
    cols = list(columns) if columns is not None else None
    lo = int(range_start) if range_start is not None else None
    hi = int(range_end) if range_end is not None else None
    range_ok = lo is not None and hi is not None and filter_column is not None
    fc = norm_col(filter_column) if filter_column else None

    if range_ok and fc:
        cols = _ensure_filter_column_in_list(cols, fc)
    cols = _ensure_filter_columns_in_list(cols, filters)

    expr_range = integer_range_expression(
        p,
        cols,
        use_pushdown=use_pushdown,
        filter_column=filter_column,
        range_start=range_start,
        range_end=range_end,
    )
    expr_filters = build_filter(filters) if filters else None
    flt = merge_dataset_expressions(expr_range, expr_filters)

    try:
        import pyarrow.dataset as ds

        dataset = ds.dataset(p, format=_parquet_read_format())
        table = _dataset_to_table(dataset, columns=cols, filter_=flt)
        df = _arrow_table_to_pandas(
            table,
            backend=be,
            dictionary_as_categorical=dictionary_as_categorical,
        )
        return _apply_unpushed_legacy_range(
            df,
            range_ok=range_ok,
            expr_range=expr_range,
            fc=fc,
            lo=lo,
            hi=hi,
        )
    except Exception:
        pass

    try:
        if cols is not None:
            df = pd.read_parquet(p, columns=cols)
        else:
            df = pd.read_parquet(p)
    except Exception:
        df = pd.read_parquet(p)

    return _finalize_parquet_dataframe(
        df,
        range_ok=range_ok,
        fc=fc,
        lo=lo,
        hi=hi,
        filters=filters,
    )


def read_csv(
    path: PathLike,
    *,
    columns: Optional[Sequence[str]] = None,
    filter_column: Optional[str] = None,
    range_start: Optional[int] = None,
    range_end: Optional[int] = None,
    filters: Optional[FilterSpec] = None,
) -> pd.DataFrame:
    """
    Read CSV with optional ``usecols``, legacy numeric range on *filter_column*, and
    optional tuple *filters* (applied in pandas; no Arrow pushdown).
    """
    p = str(path).strip()
    cols = list(columns) if columns else None
    if cols:
        try:
            df = pd.read_csv(p, usecols=cols)
        except (ValueError, KeyError, TypeError):
            df = pd.read_csv(p)
    else:
        df = pd.read_csv(p)

    lo = int(range_start) if range_start is not None else None
    hi = int(range_end) if range_end is not None else None
    if filter_column is not None and lo is not None and hi is not None:
        df = _filter_dataframe_range(df, norm_col(filter_column), lo, hi)
    if filters:
        df = apply_filters_pandas(df, filters)
    return df


def read_tabular(
    path: PathLike,
    *,
    columns: Optional[Sequence[str]] = None,
    filter_column: Optional[str] = None,
    range_start: Optional[int] = None,
    range_end: Optional[int] = None,
    use_pushdown: bool = True,
    full_read_if_no_column_match: bool = True,
    filters: Optional[FilterSpec] = None,
    backend: str = "numpy",
    dictionary_as_categorical: bool = True,
) -> pd.DataFrame:
    """
    Read ``.parquet`` or ``.csv`` from *path* (case-insensitive suffix).

    If *columns* is set, Parquet reads intersect file schema with *columns*; CSV reads
    header and uses overlapping names. When intersection is empty and
    *full_read_if_no_column_match* is True, reads all columns (then applies filters).

    Parquet-only keyword arguments: *filters*, *backend*, *dictionary_as_categorical*.
    """
    p = str(path).strip()
    low = p.lower()

    if low.endswith(".parquet"):
        sub: Optional[List[str]] = None
        if columns is not None:
            sub = intersect_columns(p, columns)
            if sub is not None and len(sub) == 0 and full_read_if_no_column_match:
                sub = None
            elif sub is not None and len(sub) == 0:
                return pd.DataFrame()

        return read_parquet(
            p,
            columns=sub,
            filter_column=filter_column,
            range_start=range_start,
            range_end=range_end,
            use_pushdown=use_pushdown,
            filters=filters,
            backend=backend,
            dictionary_as_categorical=dictionary_as_categorical,
        )

    if low.endswith(".csv"):
        sub: Optional[List[str]] = None
        if columns is not None:
            try:
                h = pd.read_csv(p, nrows=0)
                cn = {norm_col(c) for c in h.columns}
                sub = [norm_col(c) for c in columns if norm_col(c) in cn]
            except Exception:
                sub = None
            if sub is not None and len(sub) == 0 and full_read_if_no_column_match:
                sub = None
            elif sub is not None and len(sub) == 0:
                return pd.DataFrame()
        return read_csv(
            p,
            columns=sub,
            filter_column=filter_column,
            range_start=range_start,
            range_end=range_end,
            filters=filters,
        )

    raise ValueError(
        "snapparquet.read_tabular: path must end with .parquet or .csv, got {!r}".format(
            p
        )
    )


def read_tabular_many(
    paths: Sequence[PathLike],
    *,
    columns: Optional[Sequence[str]] = None,
    filter_column: Optional[str] = None,
    range_start: Optional[int] = None,
    range_end: Optional[int] = None,
    use_pushdown: bool = True,
    ignore_errors: bool = False,
    unified_parquet: bool = True,
    parallel_workers: Optional[int] = None,
    full_read_if_no_column_match: bool = True,
    filters: Optional[FilterSpec] = None,
    backend: str = "numpy",
    dictionary_as_categorical: bool = True,
) -> pd.DataFrame:
    """
    Read multiple ``.parquet`` / ``.csv`` paths and concatenate.

    * **unified_parquet** (default True): when all paths are Parquet, try one
      ``pyarrow.dataset`` scan over the whole list (often faster than opening each
      file separately; better pushdown across fragments). Falls back automatically.
    * **parallel_workers**: ``None`` = auto (parallel I/O when there are several paths),
      ``0`` or ``1`` = strict sequential single-threaded per file, or set an int cap.
    """
    raw_list = [str(raw).strip() for raw in paths if str(raw).strip()]
    if not raw_list:
        return pd.DataFrame()

    pq_only = _paths_are_all_parquet(raw_list)

    if unified_parquet and pq_only and len(raw_list) >= 1:
        uni = _try_read_parquet_unified(
            raw_list,
            columns=columns,
            filter_column=filter_column,
            range_start=range_start,
            range_end=range_end,
            use_pushdown=use_pushdown,
            filters=filters,
            backend=backend,
            dictionary_as_categorical=dictionary_as_categorical,
        )
        if uni is not None:
            return uni

    workers = parallel_workers
    if workers is None:
        workers = _default_io_workers(len(raw_list))
    else:
        workers = max(0, int(workers))

    def _one(idx_p: Tuple[int, str]) -> Tuple[int, pd.DataFrame]:
        i, p = idx_p
        df = read_tabular(
            p,
            columns=columns,
            filter_column=filter_column,
            range_start=range_start,
            range_end=range_end,
            use_pushdown=use_pushdown,
            full_read_if_no_column_match=full_read_if_no_column_match,
            filters=filters,
            backend=backend,
            dictionary_as_categorical=dictionary_as_categorical,
        )
        return i, df

    indexed = list(enumerate(raw_list))
    parts_ordered: List[Optional[pd.DataFrame]] = [None] * len(raw_list)

    if workers <= 1 or len(raw_list) == 1:
        for i, p in indexed:
            try:
                _, df = _one((i, p))
                parts_ordered[i] = df
            except Exception:
                if ignore_errors:
                    parts_ordered[i] = pd.DataFrame()
                else:
                    raise
    else:
        with ThreadPoolExecutor(max_workers=workers) as ex:
            futures = {ex.submit(_one, ip): ip[0] for ip in indexed}
            for fut in as_completed(futures):
                i = futures[fut]
                try:
                    _, df = fut.result()
                    parts_ordered[i] = df
                except Exception:
                    if ignore_errors:
                        parts_ordered[i] = pd.DataFrame()
                    else:
                        raise

    merged = [df for df in parts_ordered if df is not None and len(df) > 0]
    if not merged:
        return pd.DataFrame()
    return pd.concat(merged, ignore_index=True)


def read_parquet_many(
    paths: Sequence[PathLike],
    **kwargs,
) -> pd.DataFrame:
    """Same as :func:`read_tabular_many` (name emphasizes Parquet-only workflows)."""
    return read_tabular_many(paths, **kwargs)
