import argparse
from collections.abc import Callable, Sequence
from typing import TypeVar

from trazaeo_workflows.workflow_types import SOLANA_CLUSTERS, VERIFICATION_MODES

CommandHandler = Callable[[argparse.Namespace], None]
ParserFactory = Callable[[], argparse.ArgumentParser]
ResultT = TypeVar("ResultT")

DEFAULT_NETCDF_CHUNK_SIZE = 4096
DEFAULT_NETCDF_THREADS = 4
DEFAULT_BLAKE3_CHUNK_SIZE = 1 << 20
DEFAULT_BLAKE3_THREADS = 4
DEMO_SIGNING_KEY_HEX = "4f3edf983ac636a65a842ce7c78d9aa706d3b113bce036f9a4f5762b76f70f18"


def _parse_args(parser: argparse.ArgumentParser, argv: Sequence[str] | None) -> argparse.Namespace:
    return parser.parse_args(list(argv) if argv is not None else None)


def exit_on_command_error(action: Callable[[], ResultT]) -> ResultT:  # noqa: UP047
    try:
        return action()
    except Exception as exc:
        from examples.python_netcdf.workflow_cli_commands import CommandFailure

        if isinstance(exc, CommandFailure):
            raise SystemExit(str(exc)) from exc
        raise


def _wrap_handler(handler: CommandHandler) -> CommandHandler:
    return lambda args: exit_on_command_error(lambda: handler(args))


def _hash_netcdf_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Hash a netCDF file")
    parser.add_argument("path", help="Path to .nc file")
    parser.add_argument("--chunk", type=int, default=DEFAULT_NETCDF_CHUNK_SIZE, help="Chunk size")
    parser.add_argument("--threads", type=int, default=DEFAULT_NETCDF_THREADS, help="Thread count")
    return parser


def _reliability_check_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Compute BLAKE3 root for reliability checks")
    parser.add_argument("path", help="Path to file")
    parser.add_argument(
        "--chunk",
        type=int,
        default=DEFAULT_BLAKE3_CHUNK_SIZE,
        help="Chunk size",
    )
    parser.add_argument("--threads", type=int, default=DEFAULT_BLAKE3_THREADS, help="Thread count")
    return parser


def _publish_demo_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Build and verify a trazaeo V1 publish envelope")
    parser.add_argument(
        "--verification-mode",
        "--mode",
        dest="verification_mode",
        default="sampled",
        choices=VERIFICATION_MODES,
    )
    parser.add_argument(
        "--trust-policy-json",
        required=True,
        help="Serialized TrustPolicy JSON used to verify the publish envelope",
    )
    return parser


def _capture_parser(*, description: str) -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("--subject-id", required=True, help="Capture subject id")
    parser.add_argument(
        "--capture-actor-id",
        required=True,
        help="Attesting device or transport actor",
    )
    parser.add_argument(
        "--capture-system-id",
        required=True,
        help="Capture system, receiver pipeline, or helper-processing pipeline id",
    )
    parser.add_argument(
        "--capture-window",
        default="2026-01-01T00:00:00Z/2026-01-01T00:05:00Z",
        help="RFC3339 interval describing the capture window",
    )
    parser.add_argument("--segment-id", default="seg-1", help="Segment identifier")
    parser.add_argument(
        "--payload-text",
        default="demo-payload",
        help="UTF-8 payload text used to construct the captured segment bytes",
    )
    parser.add_argument(
        "--output-ref",
        required=True,
        help="Reference to the artifact produced by this capture event",
    )
    parser.add_argument(
        "--signer-id",
        default=None,
        help="Signer id used for the capture attestation",
    )
    parser.add_argument(
        "--signing-key-hex",
        default=DEMO_SIGNING_KEY_HEX,
        help="Hex private key for signing demo capture envelopes",
    )
    return parser


def _capture_source_parser() -> argparse.ArgumentParser:
    return _capture_parser(description="Build a demo source-device capture envelope")


def _capture_transport_parser() -> argparse.ArgumentParser:
    parser = _capture_parser(description="Build a demo transport-receipt capture envelope")
    parser.add_argument(
        "--input-ref",
        action="append",
        default=None,
        help="Reference to the upstream feed, pass, or prior transport handoff",
    )
    return parser


def _publish_solana_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Build, commit, and verify a trazaeo V1 publish envelope "
            "with a Solana proof-log provider"
        )
    )
    parser.add_argument(
        "--verification-mode",
        "--mode",
        dest="verification_mode",
        default="sampled",
        choices=VERIFICATION_MODES,
    )
    parser.add_argument("--cluster", default="solana-devnet", choices=SOLANA_CLUSTERS)
    parser.add_argument("--require-proof-log", action="store_true")
    parser.add_argument(
        "--trust-policy-json",
        required=True,
        help="Serialized TrustPolicy JSON used to verify the publish envelope",
    )
    parser.add_argument(
        "--rpc-url",
        default=None,
        help="Override the Solana RPC URL",
    )
    parser.add_argument(
        "--proof-log-keypair-path",
        default=None,
        help="Path to a funded Solana keypair JSON file for the proof-log attestor",
    )
    return parser


def _icechunk_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Convert an NC collection to Zarr/Icechunk provenance artifacts"
    )
    parser.add_argument("sources", nargs="+", help="Input netCDF file paths")
    parser.add_argument("--zarr-store", required=True, help="Output Zarr store path")
    parser.add_argument("--dataset-id", required=True, help="Dataset ID")
    parser.add_argument("--dataset-version", required=True, help="Dataset version")
    parser.add_argument("--signer-id", default="publisher", help="Attestation signer id")
    parser.add_argument(
        "--signing-key-hex",
        default=DEMO_SIGNING_KEY_HEX,
        help="Hex private key for signing test/demo envelopes",
    )
    parser.add_argument(
        "--trust-policy-json",
        required=True,
        help="Serialized TrustPolicy JSON used to verify the publish envelope",
    )
    parser.add_argument("--verification-mode", choices=VERIFICATION_MODES, default="full")
    return parser


def _icechunk_lineage_demo_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Build bootstrap and incremental Icechunk provenance bundles and verify "
            "a delivery proof for a local artifact"
        )
    )
    parser.add_argument("--bootstrap-sources", nargs="+", required=True)
    parser.add_argument("--incremental-sources", nargs="+", required=True)
    parser.add_argument("--artifact-path", required=True)
    parser.add_argument("--bootstrap-store", required=True, help="Bootstrap Icechunk store path")
    parser.add_argument(
        "--incremental-store",
        required=True,
        help="Incremental Icechunk store path",
    )
    parser.add_argument("--dataset-id", required=True, help="Dataset ID")
    parser.add_argument("--bootstrap-version", default="v1", help="Bootstrap dataset version")
    parser.add_argument(
        "--incremental-version",
        default="v2",
        help="Incremental dataset version",
    )
    parser.add_argument("--signer-id", default="publisher", help="Attestation signer id")
    parser.add_argument(
        "--signing-key-hex",
        default=DEMO_SIGNING_KEY_HEX,
        help="Hex private key for signing test/demo envelopes",
    )
    parser.add_argument(
        "--trust-policy-json",
        required=True,
        help="Serialized TrustPolicy JSON used to verify publish and delivery proofs",
    )
    parser.add_argument("--verification-mode", choices=VERIFICATION_MODES, default="full")
    return parser


def _run_hash_netcdf_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.netcdf_hashing import hash_netcdf_content_root
    from examples.python_netcdf.workflow_cli_commands import run_hash_netcdf

    run_hash_netcdf(args, hash_runner=hash_netcdf_content_root)


def _run_reliability_check_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.file_hashing import hash_file_root
    from examples.python_netcdf.workflow_cli_commands import run_reliability_check

    run_reliability_check(args, hash_runner=hash_file_root)


def _run_publish_demo_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.workflow_cli_commands import run_publish_demo

    run_publish_demo(args)


def _run_capture_source_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.workflow_cli_commands import run_capture_source

    run_capture_source(args)


def _run_capture_transport_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.workflow_cli_commands import run_capture_transport

    run_capture_transport(args)


def _run_publish_solana_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.workflow_cli_commands import run_publish_solana

    run_publish_solana(args)


def _run_icechunk_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.icechunk_conversion import convert_nc_collection_to_zarr_icechunk
    from examples.python_netcdf.workflow_cli_commands import run_icechunk

    run_icechunk(args, converter=convert_nc_collection_to_zarr_icechunk)


def _run_icechunk_lineage_demo_handler(args: argparse.Namespace) -> None:
    from examples.python_netcdf.workflow_cli_commands import run_icechunk_lineage_demo

    run_icechunk_lineage_demo(args)


def _run_single_command(
    parser_factory: ParserFactory,
    handler: CommandHandler,
    argv: Sequence[str] | None,
) -> None:
    args = _parse_args(parser_factory(), argv)
    exit_on_command_error(lambda: handler(args))


def run_hash_netcdf_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_hash_netcdf_parser, _run_hash_netcdf_handler, argv)


def run_reliability_check_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_reliability_check_parser, _run_reliability_check_handler, argv)


def run_publish_demo_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_publish_demo_parser, _run_publish_demo_handler, argv)


def run_capture_source_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_capture_source_parser, _run_capture_source_handler, argv)


def run_capture_transport_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_capture_transport_parser, _run_capture_transport_handler, argv)


def run_publish_solana_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_publish_solana_parser, _run_publish_solana_handler, argv)


def run_icechunk_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_icechunk_parser, _run_icechunk_handler, argv)


def run_icechunk_lineage_demo_cli(argv: Sequence[str] | None = None) -> None:
    _run_single_command(_icechunk_lineage_demo_parser, _run_icechunk_lineage_demo_handler, argv)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="trazaeo workflow CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    hash_parser = subparsers.add_parser(
        "hash-netcdf",
        parents=[_hash_netcdf_parser()],
        add_help=False,
    )
    hash_parser.set_defaults(handler=_wrap_handler(_run_hash_netcdf_handler))

    reliability_parser = subparsers.add_parser(
        "reliability-check",
        parents=[_reliability_check_parser()],
        add_help=False,
    )
    reliability_parser.set_defaults(handler=_wrap_handler(_run_reliability_check_handler))

    publish_demo_parser = subparsers.add_parser(
        "publish-demo",
        parents=[_publish_demo_parser()],
        add_help=False,
    )
    publish_demo_parser.set_defaults(handler=_wrap_handler(_run_publish_demo_handler))

    capture_source_parser = subparsers.add_parser(
        "capture-source",
        parents=[_capture_source_parser()],
        add_help=False,
    )
    capture_source_parser.set_defaults(handler=_wrap_handler(_run_capture_source_handler))

    capture_transport_parser = subparsers.add_parser(
        "capture-transport",
        parents=[_capture_transport_parser()],
        add_help=False,
    )
    capture_transport_parser.set_defaults(handler=_wrap_handler(_run_capture_transport_handler))

    publish_solana_parser = subparsers.add_parser(
        "publish-solana",
        parents=[_publish_solana_parser()],
        add_help=False,
    )
    publish_solana_parser.set_defaults(handler=_wrap_handler(_run_publish_solana_handler))

    icechunk_parser = subparsers.add_parser(
        "icechunk",
        parents=[_icechunk_parser()],
        add_help=False,
    )
    icechunk_parser.set_defaults(handler=_wrap_handler(_run_icechunk_handler))

    lineage_demo_parser = subparsers.add_parser(
        "icechunk-lineage-demo",
        parents=[_icechunk_lineage_demo_parser()],
        add_help=False,
    )
    lineage_demo_parser.set_defaults(handler=_wrap_handler(_run_icechunk_lineage_demo_handler))
    return parser


def main(argv: Sequence[str] | None = None) -> None:
    parser = build_parser()
    args = _parse_args(parser, argv)
    args.handler(args)
