#!/usr/bin/env python3
"""CLI helper to inspect parquet files, like `cat` for `.pq`."""

import click
import pandas as pd


@click.command(name="pqcat")
@click.argument("path", type=click.Path(exists=True))
@click.option("-n", "--head", "head", type=int, default=None, help="Show only the first N rows")
@click.option("-t", "--tail", "tail", type=int, default=None, help="Show only the last N rows")
@click.option("-c", "--columns", default=None, help="Comma-separated subset of columns to show")
@click.option("--info", is_flag=True, help="Show shape + dtypes instead of the rows")
def pqcat(path, head, tail, columns, info):
    """Dump a parquet file (or partitioned dataset dir) to stdout.

    PATH is a single .pq file or a Hive-partitioned dataset directory.
    """
    df = pd.read_parquet(path)

    if columns:
        cols = [c.strip() for c in columns.split(",") if c.strip()]
        missing = [c for c in cols if c not in df.columns]
        if missing:
            raise click.ClickException(f"Columns not found: {', '.join(missing)}")
        df = df[cols]

    if info:
        click.echo(f"path:    {path}")
        click.echo(f"shape:   {df.shape[0]:,} rows x {df.shape[1]} cols")
        click.echo("dtypes:")
        for name in df.columns:
            click.echo(f"  {name}: {df[name].dtype}")
        return

    if head is not None:
        df = df.head(head)
    elif tail is not None:
        df = df.tail(tail)

    with pd.option_context(
        "display.max_rows", None if (head is None and tail is None) else len(df),
        "display.max_columns", None,
        "display.width", None,
    ):
        click.echo(df.to_string(index=False))


if __name__ == "__main__":
    pqcat()