#!/usr/bin/env python3
"""Step 6: serve the capability under a valid licence.

With a valid leaf capability certificate, the key authority checks the chain,
releases the wrapping key, and the adapter is decrypted at load time, attached
to the base, and runs capability C at full accuracy. The plaintext adapter
exists only transiently in memory and in a gitignored scratch dir; it is never
committed.

To make the licence load-bearing, the script also shows that the exfiltrated
ciphertext from step 5, without the wrapping key, cannot be decrypted at all:
bypassing the authority with a guessed key fails AES-256-GCM authentication.
"""

import hashlib
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from demo import cryptostate
from demo.common import (
    OUTPUTS_DIR,
    evaluate,
    extract_tar_bytes,
    load_base_plus_adapter,
    load_tokenizer,
)
from demo.dataset import build_dataset

from capability_licensing import UnsealError, unseal

DECRYPT_DIR = OUTPUTS_DIR / "decrypted-adapter"


def main() -> int:
    print("=== Step 6: serve capability C under a valid licence ===\n")

    ca = cryptostate.load_ca()
    revocation = cryptostate.load_revocation()
    keys = cryptostate.load_keystore(ca, revocation)
    unit = cryptostate.load_unit()
    leaf, _leaf_key = cryptostate.load_leaf()
    exfiltrated = cryptostate.load_exfiltrated_copy()

    print(f"  presenting licence {leaf.serial} for unit '{unit.unit_id}'")
    print(f"  key state at authority: {keys.key_state(unit.unit_id).value}\n")

    # --- an attacker with only the ciphertext cannot decrypt ---
    bypass_error = None
    try:
        unseal(exfiltrated, os.urandom(32))
    except UnsealError as exc:
        bypass_error = exc
    print("  attacker holding the exfiltrated ciphertext but no key:")
    print(f"    unseal with a guessed key -> {type(bypass_error).__name__}: "
          f"{str(bypass_error)[:70]}...\n")

    # --- licensed decrypt-at-load ---
    adapter_blob = keys.request_decrypt(unit, leaf)
    print(f"  licensed decrypt succeeded: {len(adapter_blob):,} bytes recovered "
          f"(sha256 {hashlib.sha256(adapter_blob).hexdigest()[:16]}...)")

    if DECRYPT_DIR.exists():
        for p in DECRYPT_DIR.iterdir():
            p.unlink()
    extract_tar_bytes(adapter_blob, DECRYPT_DIR)
    print(f"  adapter materialised (transient, gitignored): {DECRYPT_DIR}\n")

    _, held_out = build_dataset()
    tokenizer = load_tokenizer()
    model = load_base_plus_adapter(DECRYPT_DIR)
    accuracy, correct, total, _ = evaluate(model, tokenizer, held_out)

    print(f"  LICENSED-LOAD ACCURACY ON C: {accuracy:.3f}  ({correct}/{total})")

    if accuracy < 0.80:
        print("\nRESULT: WEAK - licensed load ran but accuracy is low.")
        return 1
    print("\nRESULT: PASS - a valid licence decrypts the adapter and runs capability C.")
    return 0


if __name__ == "__main__":
    sys.exit(main())