from __future__ import annotations

import argparse
import logging
import os
import sys
import time
from pathlib import Path

from .parallel_read import RunOptions, run_parallel_read
from .metrics import format_summary

logger = logging.getLogger("universe_reads")


def _parse_args(argv: list[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Read UniVerse records using parallel sessions (uopy).")
    sub = p.add_subparsers(dest="command")

    # ── inspect sub-command ──────────────────────────────────────────────
    insp = sub.add_parser("inspect", help="Quick introspection of a UniVerse file.")
    insp.add_argument("file_name", help="UniVerse file name to inspect.")
    insp.add_argument(
        "--session-factory",
        default=None,
        help="Python callable that returns a new uopy Session.",
    )
    insp.add_argument(
        "--samples",
        type=int,
        default=3,
        help="Number of sample records to read (default: 3).",
    )
    insp.add_argument(
        "--json",
        action="store_true",
        help="Output inspect result as JSON to stdout.",
    )

    # ── read (default) ───────────────────────────────────────────────────
    read = sub.add_parser("read", help="Parallel bulk read (default when no subcommand).")
    _add_read_args(read)

    # Also accept read args directly on the root parser for backwards compat.
    _add_read_args(p)

    return p.parse_args(argv)


def _add_read_args(p: argparse.ArgumentParser) -> None:
    """Register all read-mode CLI arguments on *p*."""
    p.add_argument(
        "--file",
        dest="file_name",
        default=os.getenv("UV_FILE"),
        help="UniVerse file name to read (defaults to env UV_FILE).",
    )
    p.add_argument(
        "--session-factory",
        default=None,
        help="Python callable that returns a new uopy Session, in form 'module.sub:factory'.",
    )

    pooling = p.add_mutually_exclusive_group()
    pooling.add_argument(
        "--pooling",
        action="store_true",
        help="Pass pooling hints to the session factory (if it supports them).",
    )
    pooling.add_argument(
        "--no-pooling",
        action="store_true",
        help="Disable pooling hints to the session factory (if it supports them).",
    )

    p.add_argument("--min-pool-size", type=int, default=None)
    p.add_argument("--max-pool-size", type=int, default=None)

    p.add_argument(
        "--sessions",
        type=int,
        default=5,
        help="Number of concurrent UniVerse sessions/processes to use (default: 5).",
    )
    p.add_argument(
        "--count-only",
        action="store_true",
        help="Do not write records; just count and report throughput.",
    )
    p.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Write output files (one per worker) into this directory.",
    )
    p.add_argument(
        "--batch-size",
        type=int,
        default=500,
        help="Number of keys sent per IPC batch to each worker.",
    )
    p.add_argument(
        "--progress-every",
        type=float,
        default=5.0,
        help="Seconds between throughput logs per process.",
    )
    p.add_argument(
        "--max-records",
        type=int,
        default=None,
        help="Debug cap. Prefer leaving unset for full runs.",
    )
    p.add_argument(
        "--dry-run",
        type=int,
        default=None,
        help="If set, generate N fake keys instead of connecting to UniVerse.",
    )
    p.add_argument(
        "--retries",
        type=int,
        default=3,
        help="Number of retries per batch/key on RPC failure (default: 3).",
    )
    p.add_argument(
        "--retry-backoff",
        type=float,
        default=0.5,
        help="Base backoff in seconds between retries (default: 0.5).",
    )

    # ── new: output format ───────────────────────────────────────────────
    p.add_argument(
        "--format",
        dest="output_format",
        choices=["ndjson", "jsonl", "csv", "parquet"],
        default="ndjson",
        help="Output file format (default: ndjson).",
    )
    p.add_argument(
        "--stream",
        action="store_true",
        help="Stream records to stdout for Unix piping instead of writing files.",
    )

    # ── new: record filtering ────────────────────────────────────────────
    p.add_argument(
        "--where",
        dest="where_clause",
        default=None,
        help="UniVerse WHERE/WITH clause, e.g. \"WITH STATUS = 'ACTIVE'\".",
    )
    p.add_argument(
        "--select",
        dest="select_override",
        default=None,
        help="Full UniVerse SELECT statement override.",
    )

    # ── new: key-range partitioning ──────────────────────────────────────
    p.add_argument(
        "--partition",
        choices=["prefix", "hash"],
        default=None,
        help="Key partitioning strategy. Default is round-robin modulo.",
    )


def main(argv: list[str] | None = None) -> int:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    ns = _parse_args(sys.argv[1:] if argv is None else argv)

    # ── inspect sub-command ──────────────────────────────────────────────
    if ns.command == "inspect":
        from .inspect_file import inspect, inspect_json, print_inspect

        if not ns.session_factory:
            raise RuntimeError("Missing --session-factory for inspect")
        info = inspect(
            file_name=ns.file_name,
            session_factory_spec=ns.session_factory,
            sample_count=ns.samples,
        )
        if ns.json:
            inspect_json(info)
        else:
            print_inspect(info)
        return 0

    # ── read (default) ───────────────────────────────────────────────────
    if ns.output_dir is None and not ns.count_only and not ns.stream:
        ns.count_only = True

    if ns.dry_run is None:
        if not ns.file_name:
            raise RuntimeError("Missing --file (or set UV_FILE)")
        if not ns.session_factory:
            raise RuntimeError("Missing --session-factory (required for non-dry-run)")
    else:
        if not ns.file_name:
            ns.file_name = ""

    if ns.output_dir is not None and ns.output_dir.exists() and not ns.output_dir.is_dir():
        raise RuntimeError("--output-dir must be a directory")

    session_factory_kwargs: dict[str, object] = {}
    pooling_on: bool | None = None

    if ns.pooling:
        pooling_on = True
    elif ns.no_pooling:
        pooling_on = False

    if ns.min_pool_size is not None or ns.max_pool_size is not None:
        pooling_on = True if pooling_on is None else pooling_on

    if pooling_on is not None:
        session_factory_kwargs["pooling_on"] = pooling_on
    if ns.min_pool_size is not None:
        session_factory_kwargs["min_pool_size"] = int(ns.min_pool_size)
    if ns.max_pool_size is not None:
        session_factory_kwargs["max_pool_size"] = int(ns.max_pool_size)

    options = RunOptions(
        sessions=int(ns.sessions),
        count_only=bool(ns.count_only),
        output_dir=ns.output_dir,
        batch_size=int(ns.batch_size),
        progress_every_seconds=float(ns.progress_every),
        max_records=ns.max_records,
        dry_run=ns.dry_run,
        session_factory_kwargs=session_factory_kwargs,
        retries=int(ns.retries),
        retry_backoff=float(ns.retry_backoff),
        output_format=getattr(ns, "output_format", "ndjson") or "ndjson",
        stream=bool(getattr(ns, "stream", False)),
        where_clause=getattr(ns, "where_clause", None),
        select_override=getattr(ns, "select_override", None),
        partition=getattr(ns, "partition", None),
    )

    started = time.monotonic()
    results = run_parallel_read(
        file_name=str(ns.file_name),
        session_factory_spec=ns.session_factory,
        options=options,
    )
    elapsed = time.monotonic() - started

    total_records = sum(r.records for r in results)
    total_errors = sum(r.errors for r in results)

    logger.info(format_summary(results, wall_seconds=elapsed))

    return 0 if total_errors == 0 else 2
