"""
Fast reads: column pruning, optional integer-column range filters, PyArrow predicate pushdown.
"""

from __future__ import annotations

from pathlib import Path
from typing import List, Optional, Sequence, Union

import pandas as pd

from snapparquet.schema import column_is_integer_type, intersect_columns, norm_col

PathLike = Union[str, Path]


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 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,
) -> pd.DataFrame:
    """
    Read a Parquet file into a pandas DataFrame.

    * **Column pruning** — pass ``columns`` to decode only those fields.
    * **Range filter** — ``filter_column`` + ``range_start`` / ``range_end`` (inclusive).
      When the filter column is integer-typed in the file, uses ``pyarrow.dataset``
      predicate pushdown when possible; otherwise filters in pandas after read.
    * If *columns* is set but omits *filter_column*, the filter column is appended
      for the read so filtering can run.
    """
    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)

    use_ds = (
        bool(use_pushdown)
        and range_ok
        and fc is not None
        and (cols is None or fc in cols)
        and column_is_integer_type(p, fc)
    )

    if use_ds:
        try:
            import pyarrow.dataset as ds

            fld = ds.field(fc)
            flt = (fld >= lo) & (fld <= hi)
            dataset = ds.dataset(p, format="parquet")
            table = dataset.to_table(columns=cols, filter=flt)
            return table.to_pandas()
        except Exception:
            pass

    try:
        if cols is not None:
            df = pd.read_parquet(p, columns=cols)
            if range_ok and fc:
                df = _filter_dataframe_range(df, fc, lo, hi)
            return df
    except Exception:
        pass

    df = pd.read_parquet(p)
    if range_ok and fc:
        df = _filter_dataframe_range(df, fc, lo, hi)
    return df


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,
) -> pd.DataFrame:
    """Read CSV with optional ``usecols`` and the same range filter semantics as Parquet."""
    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)
    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,
) -> 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 range filter).
    """
    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,
        )

    if low.endswith(".csv"):
        sub = 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 not sub:
            sub = None
        return read_csv(
            p,
            columns=sub,
            filter_column=filter_column,
            range_start=range_start,
            range_end=range_end,
        )

    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,
) -> pd.DataFrame:
    """
    Read multiple ``.parquet`` / ``.csv`` paths with the same options and ``pd.concat``.

    Column pruning is applied per file via :func:`read_tabular`.
    """
    if not paths:
        return pd.DataFrame()

    parts: List[pd.DataFrame] = []
    for raw in paths:
        p = str(raw).strip()
        if not p:
            continue
        try:
            parts.append(
                read_tabular(
                    p,
                    columns=columns,
                    filter_column=filter_column,
                    range_start=range_start,
                    range_end=range_end,
                    use_pushdown=use_pushdown,
                )
            )
        except Exception:
            if ignore_errors:
                continue
            raise

    if not parts:
        return pd.DataFrame()
    return pd.concat(parts, 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)
