from __future__ import annotations

import argparse
import multiprocessing as mp
import os
import queue
import tempfile
import time

from universe_reads import run
from universe_reads.import_utils import call_factory, load_callable
from universe_reads.metrics import format_summary


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Dry-run harness for universe-reads")
    p.add_argument("--records", type=int, default=300_000, help="Number of fake records to process")
    p.add_argument("--sessions", type=int, default=5, help="Number of worker sessions/processes")
    p.add_argument("--batch-size", type=int, default=1000, help="Keys per IPC batch")
    p.add_argument("--progress-every", type=float, default=1.0, help="Seconds between progress logs")
    p.add_argument("--format", dest="output_format", choices=["ndjson", "csv", "parquet"], default="ndjson")
    p.add_argument("--partition", choices=["prefix", "hash"], default=None)
    p.add_argument("--stream", action="store_true")

    p.add_argument(
        "--smoke-session-factory",
        action="store_true",
        help=(
            "Run a multiprocessing smoke test that imports and calls a session factory by spec. "
            "This validates the 'module:callable' approach and kwargs filtering without needing UniVerse."
        ),
    )
    p.add_argument(
        "--session-factory",
        type=str,
        default="tests.test_factories:make_noarg_session",
        help=(
            "Session factory spec to smoke-test (e.g. universe_reads.session_factory:make_session). "
            "Default is an in-package dummy factory."
        ),
    )
    p.add_argument(
        "--session-factory-kw",
        type=str,
        default="tests.test_factories:make_kw_session",
        help="Second factory spec that accepts pooling kwargs (for kwargs-forwarding test).",
    )
    p.add_argument(
        "--smoke-procs",
        type=int,
        default=3,
        help="Number of processes to use for the session-factory smoke test.",
    )

    return p.parse_args()


def _smoke_worker(*, spec: str, kwargs: dict[str, object], out_q: "mp.Queue[object]") -> None:
    try:
        factory = load_callable(spec)
        session = call_factory(factory, kwargs)
        close = getattr(session, "close", None)
        if callable(close):
            close()
        out_q.put({"ok": True, "spec": spec})
    except Exception as exc:
        out_q.put({"ok": False, "spec": spec, "error": repr(exc)})


def run_session_factory_smoke_test(*, procs: int, spec_noarg: str, spec_kw: str) -> int:
    ctx = mp.get_context("spawn")
    out_q = ctx.Queue()

    # These are the same kwargs we pass through the CLI/library.
    pooling_kwargs: dict[str, object] = {
        "pooling_on": True,
        "min_pool_size": 1,
        "max_pool_size": 10,
    }

    processes: list[mp.Process] = []
    for i in range(max(1, int(procs))):
        # Alternate between no-arg factory and kw-accepting factory.
        spec = spec_kw if i % 2 else spec_noarg
        p = ctx.Process(
            target=_smoke_worker,
            kwargs={"spec": spec, "kwargs": pooling_kwargs, "out_q": out_q},
            name=f"sf-smoke-{i}",
        )
        p.daemon = True
        p.start()
        processes.append(p)

    ok = 0
    bad: list[dict[str, object]] = []

    deadline = time.monotonic() + 20.0
    remaining = len(processes)
    while remaining > 0 and time.monotonic() < deadline:
        try:
            msg = out_q.get(timeout=0.5)
        except queue.Empty:
            continue

        remaining -= 1
        if isinstance(msg, dict) and msg.get("ok") is True:
            ok += 1
        else:
            bad.append(msg if isinstance(msg, dict) else {"ok": False, "error": repr(msg)})

    for p in processes:
        p.join(timeout=2.0)

    print("\nSession factory smoke test")
    print(f"ok={ok} total={len(processes)}")
    for item in bad:
        print(f"FAIL: {item}")

    return 0 if ok == len(processes) else 2


def main() -> int:
    ns = parse_args()

    # Ensure imports work without installation:
    #   PYTHONPATH=src:. python -m tests.test_dry_run
    if ns.smoke_session_factory:
        rc = run_session_factory_smoke_test(
            procs=int(ns.smoke_procs),
            spec_noarg=str(ns.session_factory),
            spec_kw=str(ns.session_factory_kw),
        )
        if rc != 0:
            return rc

    # -------------------------------------------------------------------
    # Basic dry-run (count-only)
    # -------------------------------------------------------------------
    started = time.monotonic()
    results = run(
        file_name="",
        dry_run=int(ns.records),
        sessions=int(ns.sessions),
        batch_size=int(ns.batch_size),
        progress_every_seconds=float(ns.progress_every),
        session_factory=None,
        count_only=True,
    )
    elapsed = time.monotonic() - started

    total = sum(r.records for r in results)
    errors = sum(r.errors for r in results)

    # Use format_summary for tabular output
    print(format_summary(results, wall_seconds=elapsed))

    # Legacy per-worker printout (for comparison)
    print("Summary (legacy)")
    for r in results:
        rate = r.records / max(1e-9, r.elapsed_seconds)
        print(
            f"worker-{r.worker_id}: records={r.records} errors={r.errors} rate={rate:,.0f} rec/s elapsed={r.elapsed_seconds:,.3f}s"
        )

    print(f"TOTAL: records={total} errors={errors} elapsed={elapsed:,.3f}s")
    if elapsed > 0:
        print(f"OVERALL RATE: {total/elapsed:,.0f} rec/s")

    # -------------------------------------------------------------------
    # Dry-run with partition="hash"
    # -------------------------------------------------------------------
    print("\n--- Dry-run with partition='hash' ---")
    t0 = time.monotonic()
    results_hash = run(
        file_name="",
        dry_run=min(int(ns.records), 5000),
        sessions=int(ns.sessions),
        batch_size=int(ns.batch_size),
        progress_every_seconds=float(ns.progress_every),
        session_factory=None,
        count_only=True,
        partition="hash",
    )
    hash_elapsed = time.monotonic() - t0
    print(format_summary(results_hash, wall_seconds=hash_elapsed))
    hash_total = sum(r.records for r in results_hash)
    assert hash_total > 0, "hash-partition dry-run returned 0 records"

    # -------------------------------------------------------------------
    # Dry-run with partition="prefix"
    # -------------------------------------------------------------------
    print("\n--- Dry-run with partition='prefix' ---")
    t0 = time.monotonic()
    results_prefix = run(
        file_name="",
        dry_run=min(int(ns.records), 5000),
        sessions=int(ns.sessions),
        batch_size=int(ns.batch_size),
        progress_every_seconds=float(ns.progress_every),
        session_factory=None,
        count_only=True,
        partition="prefix",
    )
    prefix_elapsed = time.monotonic() - t0
    print(format_summary(results_prefix, wall_seconds=prefix_elapsed))
    prefix_total = sum(r.records for r in results_prefix)
    assert prefix_total > 0, "prefix-partition dry-run returned 0 records"

    # -------------------------------------------------------------------
    # Dry-run with output_format="csv" and an output directory
    # -------------------------------------------------------------------
    print("\n--- Dry-run with output_format='csv' to temp dir ---")
    with tempfile.TemporaryDirectory() as tmpdir:
        t0 = time.monotonic()
        results_csv = run(
            file_name="",
            dry_run=min(int(ns.records), 2000),
            sessions=2,
            batch_size=int(ns.batch_size),
            progress_every_seconds=float(ns.progress_every),
            session_factory=None,
            count_only=False,
            output_dir=tmpdir,
            output_format="csv",
        )
        csv_elapsed = time.monotonic() - t0
        print(format_summary(results_csv, wall_seconds=csv_elapsed))
        csv_files = [f for f in os.listdir(tmpdir) if f.endswith(".csv")]
        print(f"CSV files written: {csv_files}")
        assert len(csv_files) > 0, "Expected at least one CSV file in output dir"

    # -------------------------------------------------------------------
    # Dry-run with output_format="ndjson" and an output directory
    # -------------------------------------------------------------------
    print("\n--- Dry-run with output_format='ndjson' to temp dir ---")
    with tempfile.TemporaryDirectory() as tmpdir:
        t0 = time.monotonic()
        results_ndjson = run(
            file_name="",
            dry_run=min(int(ns.records), 2000),
            sessions=2,
            batch_size=int(ns.batch_size),
            progress_every_seconds=float(ns.progress_every),
            session_factory=None,
            count_only=False,
            output_dir=tmpdir,
            output_format="ndjson",
        )
        ndjson_elapsed = time.monotonic() - t0
        print(format_summary(results_ndjson, wall_seconds=ndjson_elapsed))
        ndjson_files = [f for f in os.listdir(tmpdir) if f.endswith(".ndjson")]
        print(f"NDJSON files written: {ndjson_files}")
        assert len(ndjson_files) > 0, "Expected at least one NDJSON file in output dir"

    # Print a suggested command line for convenience
    print("\nRun command:")
    print(f"  PYTHONPATH=src:. python -m tests.test_dry_run --records {ns.records} --sessions {ns.sessions}")
    print("\nSession factory smoke test command:")
    print(
        "  PYTHONPATH=src:. python -m tests.test_dry_run --smoke-session-factory "
        "--session-factory tests.test_factories:make_noarg_session "
        "--session-factory-kw tests.test_factories:make_kw_session"
    )

    return 0 if errors == 0 else 2


if __name__ == "__main__":
    raise SystemExit(main())
