"""
Pipeline — Declarative data transformation pipeline with caching and logging.
"""

from __future__ import annotations

import time
import logging
from dataclasses import dataclass, field
from typing import Any, Callable

import pandas as pd

from bizanalyticsai.core.dataframe import BizFrame

logger = logging.getLogger(__name__)


@dataclass
class StepResult:
    name:       str
    duration_s: float
    input_shape:  tuple[int, int]
    output_shape: tuple[int, int]
    error:      str | None = None

    @property
    def ok(self) -> bool:
        return self.error is None

    def __repr__(self) -> str:
        status = "✓" if self.ok else "✗"
        return (
            f"{status} {self.name} | {self.input_shape} → {self.output_shape} "
            f"| {self.duration_s:.3f}s"
            + (f" | ERROR: {self.error}" if self.error else "")
        )


class Pipeline:
    """
    Fluent data pipeline for composing transformations.

    Example::

        pipe = (
            Pipeline("sales_etl")
            .read(bai.read_csv, "sales.csv")
            .step("clean",    bai.clean)
            .step("normalize", lambda df: bai.normalize(df, cols=["revenue"]))
            .step("outliers", bai.remove_outliers)
            .step("segment",  lambda df: df.segment("revenue", bins=4))
        )

        result = pipe.run()
        print(pipe.report())
    """

    def __init__(self, name: str = "pipeline", verbose: bool = True) -> None:
        self.name     = name
        self.verbose  = verbose
        self._steps:  list[tuple[str, Callable]] = []
        self._results: list[StepResult] = []
        self._data:   BizFrame | None = None

    # ── Builder ───────────────────────────────────────────────────────────────
    def read(self, reader: Callable, *args, **kwargs) -> Pipeline:
        """Add a data-loading step as the first step."""
        def _load(_: Any) -> BizFrame:
            raw = reader(*args, **kwargs)
            return BizFrame(raw) if not isinstance(raw, BizFrame) else raw
        self._steps.insert(0, ("load", _load))
        return self

    def step(self, name: str, fn: Callable, **kwargs) -> Pipeline:
        """Add a named transform step."""
        if kwargs:
            fn = lambda df, _fn=fn, _kw=kwargs: _fn(df, **_kw)
        self._steps.append((name, fn))
        return self

    def filter(self, condition: str | Callable) -> Pipeline:
        """Add a filter step (SQL-style string or callable)."""
        if isinstance(condition, str):
            fn = lambda df: BizFrame(df.query(condition))
        else:
            fn = lambda df: BizFrame(df[condition(df)])
        return self.step(f"filter({condition if isinstance(condition, str) else condition.__name__})", fn)

    def select(self, *cols: str) -> Pipeline:
        """Keep only the specified columns."""
        return self.step(f"select({list(cols)})", lambda df: BizFrame(df[list(cols)]))

    def rename(self, mapping: dict[str, str]) -> Pipeline:
        return self.step("rename", lambda df: BizFrame(df.rename(columns=mapping)))

    def cast(self, mapping: dict[str, str]) -> Pipeline:
        """Cast columns to given dtypes."""
        def _cast(df: BizFrame) -> BizFrame:
            return BizFrame(df.astype(mapping))
        return self.step(f"cast({list(mapping.keys())})", _cast)

    def sort(self, col: str, ascending: bool = True) -> Pipeline:
        return self.step(f"sort({col})", lambda df: BizFrame(df.sort_values(col, ascending=ascending)))

    def limit(self, n: int) -> Pipeline:
        return self.step(f"limit({n})", lambda df: BizFrame(df.head(n)))

    def apply(self, fn: Callable, col: str, out_col: str | None = None) -> Pipeline:
        """Apply a function element-wise to a column."""
        def _apply(df: BizFrame) -> BizFrame:
            out = df.copy()
            out[out_col or col] = out[col].apply(fn)
            return BizFrame(out)
        return self.step(f"apply({col}→{out_col or col})", _apply)

    def join(self, other: pd.DataFrame, on: str | list[str], how: str = "left") -> Pipeline:
        def _join(df: BizFrame) -> BizFrame:
            return BizFrame(df.merge(other, on=on, how=how))
        return self.step(f"join(on={on}, how={how})", _join)

    # ── Execution ─────────────────────────────────────────────────────────────
    def run(self, data: pd.DataFrame | None = None) -> BizFrame:
        """Execute the pipeline and return the final BizFrame."""
        self._results = []
        current: Any = data

        for name, fn in self._steps:
            t0 = time.perf_counter()
            in_shape = current.shape if hasattr(current, "shape") else (0, 0)
            try:
                current = fn(current)
                if not isinstance(current, BizFrame):
                    current = BizFrame(current)
                out_shape = current.shape
                err = None
            except Exception as exc:
                logger.error("Pipeline step '%s' failed: %s", name, exc)
                err = str(exc)
                out_shape = in_shape

            result = StepResult(
                name=name,
                duration_s=round(time.perf_counter() - t0, 4),
                input_shape=in_shape,
                output_shape=out_shape,
                error=err,
            )
            self._results.append(result)
            if self.verbose:
                print(repr(result))
            if err:
                raise RuntimeError(f"Pipeline failed at step '{name}': {err}") from None

        self._data = current
        return current

    def report(self) -> str:
        """Return a text report of the pipeline execution."""
        lines = [f"\n{'=' * 60}", f"Pipeline: {self.name}", "=" * 60]
        total = sum(r.duration_s for r in self._results)
        for r in self._results:
            lines.append(repr(r))
        lines += ["=" * 60, f"Total steps: {len(self._results)} | Total time: {total:.3f}s", "=" * 60]
        return "\n".join(lines)

    def __repr__(self) -> str:
        steps = " → ".join(n for n, _ in self._steps)
        return f"<Pipeline '{self.name}': {steps}>"
