"""Dataset Export: snapshot features + labels to Parquet."""
from __future__ import annotations

import argparse
import asyncio
import json

from libs.common.config import get_settings
from libs.common.logging import configure_logging, get_logger
from libs.db.session import get_session
from libs.export.snapshot_export import export_dataset_snapshot

logger = get_logger(__name__)


async def run_dataset_export(
    snapshot_id: str | None,
    split_policy: str,
    output_dir: str,
) -> dict:
    async with get_session() as session:
        manifest = await export_dataset_snapshot(
            session=session,
            snapshot_id=snapshot_id,
            split_policy=split_policy,
            output_dir=output_dir,
        )
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(description="Dataset Export")
    parser.add_argument("--snapshot-id", default=None, help="Snapshot ID (UUID, auto-generated if omitted)")
    parser.add_argument(
        "--split-policy",
        default="temporal_70_15_15",
        help="Split policy string (default: temporal_70_15_15)",
    )
    parser.add_argument(
        "--output-dir",
        default="./data/datasets/snapshots",
        help="Output directory for Parquet files",
    )
    parser.add_argument("--json", action="store_true", help="Print manifest JSON to stdout")
    args = parser.parse_args()

    settings = get_settings()
    configure_logging(settings.log_level)

    manifest = asyncio.run(
        run_dataset_export(
            snapshot_id=args.snapshot_id,
            split_policy=args.split_policy,
            output_dir=args.output_dir,
        )
    )

    if args.json:
        print(json.dumps(manifest, indent=2))
    else:
        print(f"Snapshot exported: {manifest['snapshot_id']}")
        print(f"  Output: {manifest['output_dir']}")
        print(f"  Total rows: {manifest['total_rows']}")
        for split, count in manifest["row_counts"].items():
            print(f"  {split}: {count} rows")


if __name__ == "__main__":
    main()