"""Top-level Aliro reader transaction orchestrator.

Runs SELECT EXPEDITED -> AUTH0 -> AUTH1 against an abstract APDU transport,
decrypts the AUTH1 response, and verifies the UD signature against the
configured credential_PubK. Decoupled from pyscard so unit tests can drive
a mocked/in-process card.
"""

import secrets
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec

from aliro_harness.personalizer.credentials import extract_pub_xy
from aliro_harness.reader.auth0 import build_auth0_apdu, build_auth0_data
from aliro_harness.reader.auth1 import (
    build_auth1_apdu,
    build_auth1_data,
    build_table_812,
    build_table_813,
    sign_table_812_raw,
)
from aliro_harness.reader.auth1_response import (
    decrypt_auth1_response,
    verify_ud_signature,
)
from aliro_harness.reader.crypto import derive_kdh
from aliro_harness.reader.key_derivation import (
    OFF_EXPEDITED_SK_DEVICE,
    OFF_STEP_UP_SK,
    build_salt_volatile,
    derive_expedited_session_keys,
)
from aliro_harness.reader.tlv import find_top_level

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

Transmit = Callable[[bytes], tuple[bytes, int]]
SW_OK = 0x9000


@dataclass(frozen=True)
class TrustBundle:
    """Trust artifacts the reader needs for a transaction. Mirrors what
    personalizer.TrustArtifacts loads but adds reader_priv (which the
    personalizer doesn't need, since it never signs)."""

    reader_priv: ec.EllipticCurvePrivateKey
    reader_long_term_pub_x: bytes  # 32B (used in salt_volatile)
    reader_group_id: bytes  # 16B
    reader_group_sub_id: bytes  # 16B
    credential_long_term_pub_uncompressed: bytes  # 65B (for verify_ud_signature)
    credential_long_term_pub_x: bytes  # 32B (for salt_volatile)

    @classmethod
    def synthesize(
        cls,
        *,
        reader_priv: ec.EllipticCurvePrivateKey,
        credential_pub: ec.EllipticCurvePublicKey,
        reader_group_id: bytes | None = None,
        reader_group_sub_id: bytes | None = None,
    ) -> "TrustBundle":
        """Convenience for tests: synthesize a TrustBundle from in-process keys."""
        reader_pub_xy = extract_pub_xy(reader_priv.public_key())  # 64B (x||y)
        reader_pub_x = reader_pub_xy[:32]

        credential_pub_xy = extract_pub_xy(credential_pub)  # 64B
        credential_pub_x = credential_pub_xy[:32]
        credential_pub_uncompressed = b"\x04" + credential_pub_xy

        if reader_group_id is None:
            reader_group_id = secrets.token_bytes(16)
        if reader_group_sub_id is None:
            reader_group_sub_id = secrets.token_bytes(16)

        if len(reader_group_id) != 16:
            raise ValueError("reader_group_id must be 16B")
        if len(reader_group_sub_id) != 16:
            raise ValueError("reader_group_sub_id must be 16B")

        return cls(
            reader_priv=reader_priv,
            reader_long_term_pub_x=reader_pub_x,
            reader_group_id=reader_group_id,
            reader_group_sub_id=reader_group_sub_id,
            credential_long_term_pub_uncompressed=credential_pub_uncompressed,
            credential_long_term_pub_x=credential_pub_x,
        )

    @classmethod
    def from_trust_dir(cls, trust_dir: Path) -> "TrustBundle":
        """Loads from ``aliro-trustgen init`` output:
          - reader.pem -> reader_priv
          - access_credential.pem -> credential public key
          - reader_group_id.bin -> reader_group_id
          - reader_group_sub_id.bin -> reader_group_sub_id
        """
        reader_pem = (trust_dir / "reader.pem").read_bytes()
        cred_pem = (trust_dir / "access_credential.pem").read_bytes()
        reader_group_id = (trust_dir / "reader_group_id.bin").read_bytes()
        reader_group_sub_id = (trust_dir / "reader_group_sub_id.bin").read_bytes()

        reader_priv = serialization.load_pem_private_key(reader_pem, password=None)
        cred_priv = serialization.load_pem_private_key(cred_pem, password=None)

        return cls.synthesize(
            reader_priv=reader_priv,
            credential_pub=cred_priv.public_key(),
            reader_group_id=reader_group_id,
            reader_group_sub_id=reader_group_sub_id,
        )


@dataclass(frozen=True)
class TransactionResult:
    ok: bool
    plaintext: bytes | None = None  # decrypted Table 8-11 bytes
    credential_pub: bytes | None = None  # 0x5A value (65B uncompressed) when present
    ud_sig: bytes | None = None  # 0x9E value (64B raw r||s)
    signaling_bitmap: bytes | None = None  # 0x5E value (2B)
    step_up_sk: bytes | None = None  # 32B StepUpSK from derived_keys_volatile[64..96)
    error: str | None = None  # populated on failure
    # Wall-clock per-APDU latencies in milliseconds. Populated as each step
    # runs, so a partial result on failure still surfaces what we measured
    # before the failing step. Keys: "select", "auth0", "auth1".
    latencies_ms: dict[str, float] = field(default_factory=dict)


def run_aliro_transaction(
    *, transmit: Transmit, bundle: TrustBundle
) -> TransactionResult:
    """Orchestrate one Aliro AUTH0+AUTH1 transaction.

    Steps:
      1. SELECT EXPEDITED
      2. AUTH0 -- generate ephemeral, send, parse 0x86 (credential_ePubK)
      3. Derive Kdh + ExpeditedSKDevice
      4. Build + sign Table 8-12 -> AUTH1 command
      5. Decrypt response, extract 0x5A/0x9E/0x5E TLVs
      6. Verify UD signature over Table 8-13 against credential_PubK

    Any non-9000 SW or signature mismatch returns ok=False with ``error`` set.
    """
    # Collect per-APDU wall-clock latencies; mutated in place and threaded
    # through every TransactionResult so partial timings surface on failure.
    latencies: dict[str, float] = {}

    # SELECT EXPEDITED
    select_apdu = (
        bytes([0x00, 0xA4, 0x04, 0x00, len(EXPEDITED_AID)]) + EXPEDITED_AID
    )
    _t = time.perf_counter()
    _, sw = transmit(select_apdu)
    latencies["select"] = (time.perf_counter() - _t) * 1000.0
    if sw != SW_OK:
        return TransactionResult(
            ok=False,
            error=f"SELECT EXPEDITED failed: SW=0x{sw:04X}",
            latencies_ms=latencies,
        )

    # Generate reader ephemeral + transaction state
    reader_ephem = ec.generate_private_key(ec.SECP256R1())
    reader_ephem_pub_uncompressed = reader_ephem.public_key().public_bytes(
        encoding=serialization.Encoding.X962,
        format=serialization.PublicFormat.UncompressedPoint,
    )
    reader_ephem_pub_x = reader_ephem_pub_uncompressed[1:33]
    txn_id = secrets.token_bytes(16)

    # AUTH0
    auth0_data = build_auth0_data(
        reader_ephem_pub_uncompressed=reader_ephem_pub_uncompressed,
        reader_group_id=bundle.reader_group_id,
        reader_group_sub_id=bundle.reader_group_sub_id,
        transaction_id=txn_id,
        command_parameters=0x00,
        authentication_policy=0x00,
    )
    _t = time.perf_counter()
    auth0_resp_data, sw = transmit(build_auth0_apdu(auth0_data))
    latencies["auth0"] = (time.perf_counter() - _t) * 1000.0
    if sw != SW_OK:
        return TransactionResult(
            ok=False,
            error=f"AUTH0 failed: SW=0x{sw:04X}",
            latencies_ms=latencies,
        )

    cred_ephem_pub_uncompressed = find_top_level(auth0_resp_data, 0x86)
    if (
        cred_ephem_pub_uncompressed is None
        or len(cred_ephem_pub_uncompressed) != 65
    ):
        return TransactionResult(
            ok=False,
            error="AUTH0 response missing or malformed 0x86 credential_ePubK",
            latencies_ms=latencies,
        )
    cred_ephem_pub_x = cred_ephem_pub_uncompressed[1:33]

    # Key derivation
    kdh = derive_kdh(reader_ephem, cred_ephem_pub_uncompressed, txn_id)
    # flag = command_parameters || authentication_policy per §8.3.1.13.
    # command_parameters here is the AUTH1 command's value (0x01 for
    # "request credential_PubK in response"), not AUTH0's. authentication_policy
    # only appears in AUTH0; v1 hardcodes 0x00 (no policy enforced).
    # See the AUTH1 build_auth1_data call below -- same value must round-trip.
    salt_volatile = build_salt_volatile(
        reader_long_term_pub_x=bundle.reader_long_term_pub_x,
        reader_group_id=bundle.reader_group_id,
        reader_group_sub_id=bundle.reader_group_sub_id,
        reader_ephem_pub_x=reader_ephem_pub_x,
        transaction_id=txn_id,
        command_parameters=0x01,
        authentication_policy=0x00,
    )
    derived_keys = derive_expedited_session_keys(
        kdh, salt_volatile, cred_ephem_pub_x
    )
    sk_device = derived_keys[OFF_EXPEDITED_SK_DEVICE : OFF_EXPEDITED_SK_DEVICE + 32]
    step_up_sk = derived_keys[OFF_STEP_UP_SK : OFF_STEP_UP_SK + 32]

    # AUTH1
    reader_id_32 = bundle.reader_group_id + bundle.reader_group_sub_id
    table_812 = build_table_812(
        reader_id_32=reader_id_32,
        credential_ephem_pub_x=cred_ephem_pub_x,
        reader_ephem_pub_x=reader_ephem_pub_x,
        transaction_id=txn_id,
    )
    raw_sig = sign_table_812_raw(table_812, bundle.reader_priv)
    auth1_data = build_auth1_data(raw_sig, command_parameters=0x01)
    _t = time.perf_counter()
    auth1_resp_data, sw = transmit(build_auth1_apdu(auth1_data))
    latencies["auth1"] = (time.perf_counter() - _t) * 1000.0
    if sw != SW_OK:
        return TransactionResult(
            ok=False,
            error=f"AUTH1 failed: SW=0x{sw:04X}",
            latencies_ms=latencies,
        )

    # Decrypt + verify
    try:
        plaintext = decrypt_auth1_response(sk_device, auth1_resp_data)
    except Exception as e:
        return TransactionResult(
            ok=False,
            error=f"AUTH1 response decrypt failed: {e}",
            latencies_ms=latencies,
        )

    cred_pub_5a = find_top_level(plaintext, 0x5A)
    ud_sig = find_top_level(plaintext, 0x9E)
    bitmap = find_top_level(plaintext, 0x5E)

    if ud_sig is None or len(ud_sig) != 64:
        return TransactionResult(
            ok=False,
            error="Table 8-11 missing 0x9E UD signature",
            plaintext=plaintext,
            latencies_ms=latencies,
        )
    if bitmap is None:
        return TransactionResult(
            ok=False,
            error="Table 8-11 missing 0x5E signaling_bitmap",
            plaintext=plaintext,
            latencies_ms=latencies,
        )

    table_813 = build_table_813(
        reader_id_32=reader_id_32,
        credential_ephem_pub_x=cred_ephem_pub_x,
        reader_ephem_pub_x=reader_ephem_pub_x,
        transaction_id=txn_id,
    )
    ud_pub = (
        cred_pub_5a
        if cred_pub_5a
        else bundle.credential_long_term_pub_uncompressed
    )
    if not verify_ud_signature(ud_pub, table_813, ud_sig):
        return TransactionResult(
            ok=False,
            error="UD signature verification failed",
            plaintext=plaintext,
            credential_pub=cred_pub_5a,
            ud_sig=ud_sig,
            signaling_bitmap=bitmap,
            latencies_ms=latencies,
        )

    return TransactionResult(
        ok=True,
        plaintext=plaintext,
        credential_pub=cred_pub_5a,
        ud_sig=ud_sig,
        signaling_bitmap=bitmap,
        step_up_sk=step_up_sk,
        latencies_ms=latencies,
    )
