"""Pipeline orchestration for data curation workflows."""

from __future__ import annotations

import time
from pathlib import Path
from typing import Any

import yaml
from rich.console import Console
from rich.table import Table

from winnow.core.filters import Filter
from winnow.core.registry import FILTER_REGISTRY

console = Console()


class Pipeline:
    """A composable data curation pipeline.

    Example::

        pipeline = Pipeline("my-pipeline")
        pipeline.add_filter(LengthFilter(min_length=100))
        pipeline.add_filter(WhitespaceFilter(max_whitespace_ratio=0.4))
        pipeline.add_filter(MinHashDedup(threshold=0.8))
        result = pipeline.run("data/raw.jsonl", "data/curated.parquet")
    """

    def __init__(self, name: str = "default") -> None:
        self.name = name
        self.filters: list[Filter] = []

    def add_filter(self, f: Filter) -> Pipeline:
        """Add a filter to the pipeline. Returns self for chaining."""
        self.filters.append(f)
        return self

    @classmethod
    def from_config(cls, config_path: str | Path) -> Pipeline:
        """Load a pipeline from a YAML config file.

        Example YAML::

            name: my-pipeline
            filters:
              - LengthFilter:
                  min_length: 100
                  max_length: 100000
              - WhitespaceFilter:
                  max_whitespace_ratio: 0.4
              - RepetitionFilter:
                  max_repetition_ratio: 0.3
              - ExactDedup: {}
        """
        with open(config_path) as f:
            config = yaml.safe_load(f)

        pipeline = cls(name=config.get("name", "default"))

        for filter_spec in config.get("filters", []):
            if isinstance(filter_spec, str):
                # Simple filter name with no args: "- ExactDedup"
                filter_cls = FILTER_REGISTRY[filter_spec]
                pipeline.add_filter(filter_cls())
            elif isinstance(filter_spec, dict):
                # Filter with args: "- LengthFilter: {min_length: 100}"
                for filter_name, filter_args in filter_spec.items():
                    filter_cls = FILTER_REGISTRY[filter_name]
                    if filter_args:
                        pipeline.add_filter(filter_cls(**filter_args))
                    else:
                        pipeline.add_filter(filter_cls())

        return pipeline

    def run(
        self,
        input_path: str,
        output_path: str | None = None,
        quiet: bool = False,
    ) -> dict[str, Any]:
        """Execute the pipeline on input data.

        Args:
            input_path: Path to input dataset (JSONL, Parquet, CSV, or HuggingFace dataset ID).
            output_path: Path for curated output. Format inferred from extension.
                         If None, runs pipeline and returns stats without writing.
            quiet: Suppress progress output.

        Returns:
            Dictionary with pipeline execution statistics.
        """
        from tqdm import tqdm

        from winnow.io.readers import auto_read
        from winnow.io.writers import write_jsonl, write_parquet

        start_time = time.time()

        # Track per-filter stats
        filter_stats: dict[str, int] = {f.name(): 0 for f in self.filters}
        total_read = 0
        total_kept = 0

        # Collect output if writing
        output_buffer: list[dict[str, Any]] = []

        if not quiet:
            console.print(f"[bold blue]Winnow[/] — Running pipeline [bold]{self.name}[/]")
            console.print(f"  Input:   {input_path}")
            console.print(f"  Output:  {output_path or '(dry run)'}")
            console.print(f"  Filters: {len(self.filters)}")
            console.print()

        reader = auto_read(input_path)
        progress = tqdm(reader, desc="Processing", unit=" docs", disable=quiet)

        for sample in progress:
            total_read += 1
            kept = True

            for f in self.filters:
                if not f.filter(sample):
                    filter_stats[f.name()] += 1
                    kept = False
                    break

            if kept:
                total_kept += 1
                if output_path is not None:
                    output_buffer.append(sample)

        progress.close()
        elapsed = time.time() - start_time

        # Write output
        if output_path is not None and output_buffer:
            out = Path(output_path)
            out.parent.mkdir(parents=True, exist_ok=True)

            if out.suffix == ".parquet":
                write_parquet(output_buffer, out)
            else:
                write_jsonl(output_buffer, out)

        # Build stats
        total_removed = total_read - total_kept
        removal_rate = (total_removed / total_read * 100) if total_read > 0 else 0.0
        docs_per_sec = total_read / elapsed if elapsed > 0 else 0.0

        stats: dict[str, Any] = {
            "pipeline": self.name,
            "input": input_path,
            "output": output_path,
            "total_read": total_read,
            "total_kept": total_kept,
            "total_removed": total_removed,
            "removal_rate_pct": round(removal_rate, 2),
            "filters_applied": len(self.filters),
            "per_filter_removed": filter_stats,
            "elapsed_seconds": round(elapsed, 2),
            "docs_per_second": round(docs_per_sec, 1),
        }

        if not quiet:
            self._print_stats(stats)

        return stats

    def _print_stats(self, stats: dict[str, Any]) -> None:
        """Print a summary table of pipeline results."""
        console.print()

        table = Table(title="Pipeline Results", show_header=True)
        table.add_column("Metric", style="bold")
        table.add_column("Value", justify="right")

        table.add_row("Documents read", f"{stats['total_read']:,}")
        table.add_row("Documents kept", f"[green]{stats['total_kept']:,}[/]")
        table.add_row("Documents removed", f"[red]{stats['total_removed']:,}[/]")
        table.add_row("Removal rate", f"{stats['removal_rate_pct']}%")
        table.add_row("Elapsed", f"{stats['elapsed_seconds']}s")
        table.add_row("Throughput", f"{stats['docs_per_second']:,.0f} docs/s")

        console.print(table)

        # Per-filter breakdown
        if stats["per_filter_removed"]:
            console.print()
            ft = Table(title="Per-Filter Removals", show_header=True)
            ft.add_column("Filter", style="bold")
            ft.add_column("Removed", justify="right")

            for fname, count in stats["per_filter_removed"].items():
                ft.add_row(fname, f"{count:,}")

            console.print(ft)
