"""``aliro-bench-profile`` -- per-primitive latency profiling of the Aliro
applet on real hardware via PC/SC.

Selects the Aliro Expedited AID, then drives the four INS_DIAG_* commands
the applet exposes for profiling. Each command runs an operation N times
against hardcoded card-side test vectors. We measure two N values per
operation and compute per-op cost as the slope, which factors out the
fixed APDU round-trip and any one-time setup (Signature.init etc.).

Output: a breakdown of per-op cost in milliseconds for HMAC, ECDH, ECDSA
sign, and full AES-256-GCM encrypt of 137 bytes -- the same primitives
AUTH1 invokes. Plus a reconstructed AUTH1 budget computed from the
measured per-op costs, so we can compare against the bench-test wall-clock.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from pathlib import Path

import click

# Aliro Expedited AID -- must match AliroAids.EXPEDITED in the applet.
EXPEDITED_AID = bytes.fromhex("A000000909ACCE5501")

INS_DIAG_HMAC = 0xD0
INS_DIAG_ECDH = 0xD1
INS_DIAG_ECDSA_SIGN = 0xD2
INS_DIAG_GCM = 0xD3

# AUTH1 invokes each primitive this many times. From the applet source:
#   ECDSA verify (reader sig over Table 8-12): 1   -- modeled as ECDSA sign
#   ECDH (Kdh derivation):                     1
#   HMAC (HKDF Extract + Expand for 160B):     6   -- 1 Extract + 5 Expand
#   ECDSA sign (UD signature over Table 8-13): 1
#   AES-GCM encrypt of 137B Table 8-11:        1
AUTH1_OP_COUNTS = {
    "ecdsa_verify": 1,  # use ECDSA-sign cost as proxy; same algorithm
    "ecdh": 1,
    "hmac": 6,
    "ecdsa_sign": 1,
    "gcm_137": 1,
}


@dataclass
class OpResult:
    name: str
    per_op_ms: float
    n_low: int
    t_low_ms: float
    n_high: int
    t_high_ms: float


@click.command()
@click.option(
    "--reader",
    "reader_index",
    type=int,
    default=0,
    show_default=True,
    help="Index into PC/SC reader list",
)
@click.option(
    "--list-readers",
    is_flag=True,
    help="List PC/SC readers and exit",
)
@click.option(
    "--n-low",
    type=int,
    default=4,
    show_default=True,
    help="Low iteration count for slope measurement",
)
@click.option(
    "--n-high",
    type=int,
    default=16,
    show_default=True,
    help="High iteration count for slope measurement",
)
def main(reader_index: int, list_readers: bool, n_low: int, n_high: int) -> None:
    """Profile AUTH1 sub-operations on the Aliro applet via PC/SC.

    Card must have AliroApplet installed (personalization not required for
    diagnostics -- they use hardcoded test vectors).
    """
    from smartcard.System import readers

    available = readers()
    if list_readers:
        for i, r in enumerate(available):
            click.echo(f"[{i}] {r}")
        return

    if not available:
        raise click.ClickException("No PC/SC readers found. Is pcscd running?")
    if not (0 <= reader_index < len(available)):
        raise click.ClickException(f"--reader {reader_index} out of range")
    if n_low >= n_high:
        raise click.ClickException("--n-low must be smaller than --n-high")
    if not (1 <= n_low <= 255 and 1 <= n_high <= 255):
        raise click.ClickException("N values must be in [1, 255]")

    reader = available[reader_index]
    click.echo(f"Using reader: {reader}")
    connection = reader.createConnection()
    connection.connect()
    try:
        click.echo(f"Connected. ATR: {bytes(connection.getATR()).hex()}")

        def transmit(apdu: bytes) -> tuple[bytes, int]:
            data, sw1, sw2 = connection.transmit(list(apdu))
            return bytes(data), (sw1 << 8) | sw2

        # SELECT Aliro Expedited AID
        select_apdu = bytes([0x00, 0xA4, 0x04, 0x00, len(EXPEDITED_AID)]) + EXPEDITED_AID
        _, sw = transmit(select_apdu)
        if sw != 0x9000:
            raise click.ClickException(f"SELECT failed: SW=0x{sw:04X}")
        click.echo(f"SELECT Aliro Expedited OK (SW=0x{sw:04X})\n")

        ops = [
            ("hmac", INS_DIAG_HMAC, "HMAC-SHA-256 (32B key, 64B msg)"),
            ("ecdh", INS_DIAG_ECDH, "ECDH P-256 (native)"),
            ("ecdsa_sign", INS_DIAG_ECDSA_SIGN, "ECDSA-SHA-256 sign (native)"),
            ("gcm_137", INS_DIAG_GCM, "AES-256-GCM encrypt 137B (userland)"),
        ]

        results: list[OpResult] = []
        for key, ins, label in ops:
            r = _profile_one(transmit, ins, n_low, n_high, key)
            results.append(r)
            click.echo(
                f"  {label:<48} "
                f"N={n_low:>3}: {r.t_low_ms:>7.1f} ms  "
                f"N={n_high:>3}: {r.t_high_ms:>7.1f} ms  "
                f"-> {r.per_op_ms:>6.1f} ms/op"
            )

        click.echo()
        _print_auth1_reconstruction(results)
    finally:
        connection.disconnect()


def _profile_one(
    transmit, ins: int, n_low: int, n_high: int, name: str
) -> OpResult:
    """Send the diagnostic INS twice with different iteration counts and
    return per-op cost computed from the slope. The slope cancels out APDU
    overhead and any one-time setup that doesn't depend on N."""
    # APDU: CLA=80 INS=ins P1=00 P2=00 Lc=01 [N]
    def measure(n: int) -> float:
        apdu = bytes([0x80, ins, 0x00, 0x00, 0x01, n & 0xFF])
        t = time.perf_counter()
        _, sw = transmit(apdu)
        dt = (time.perf_counter() - t) * 1000.0
        if sw != 0x9000:
            raise click.ClickException(
                f"Diag INS 0x{ins:02X} (N={n}) failed: SW=0x{sw:04X}"
            )
        return dt

    t_low = measure(n_low)
    t_high = measure(n_high)
    per_op = (t_high - t_low) / (n_high - n_low)
    return OpResult(
        name=name,
        per_op_ms=per_op,
        n_low=n_low,
        t_low_ms=t_low,
        n_high=n_high,
        t_high_ms=t_high,
    )


def _print_auth1_reconstruction(results: list[OpResult]) -> None:
    """Reconstruct an AUTH1 budget from the measured per-op costs. Compares
    the sum to the bench-test wall-clock so we can see if the primitives
    fully account for the AUTH1 latency (or if there's additional overhead
    we haven't profiled, like TLV parsing, key loading, etc.)."""
    by_name = {r.name: r.per_op_ms for r in results}
    # Use ecdsa_sign cost as ECDSA verify proxy (same algorithm, similar work).
    by_name["ecdsa_verify"] = by_name["ecdsa_sign"]

    click.echo("AUTH1 reconstructed budget (using AUTH1_OP_COUNTS):")
    total = 0.0
    for op, count in AUTH1_OP_COUNTS.items():
        per_op = by_name.get(op, 0.0)
        sub_total = per_op * count
        total += sub_total
        click.echo(
            f"  {op:<14} x {count:>2}  ->  {sub_total:>7.1f} ms  "
            f"({per_op:.1f} ms/op)"
        )
    click.echo(f"  {'estimated_total':<19}  {total:>7.1f} ms")
    click.echo(
        f"\nCompare against `aliro-bench-test` AUTH1 wall-clock. Gap "
        f"= TLV parsing + key load + apdu I/O on card."
    )


if __name__ == "__main__":
    main()
