#!/usr/bin/env python3
import argparse
import base64
import copy
import json
import subprocess
import tempfile
from pathlib import Path

from cezen_license import utcnow_iso


def canonical_payload(payload):
    clean = copy.deepcopy(payload)
    clean.pop("signature", None)
    clean.pop("signature_status", None)
    return json.dumps(clean, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode()


def main():
    parser = argparse.ArgumentParser(description="Sign a Nexus One AI license or override artifact")
    parser.add_argument("--input", required=True, help="Unsigned JSON artifact")
    parser.add_argument("--private-key", required=True, help="PEM private key path")
    parser.add_argument("--output", required=True, help="Signed JSON output path")
    parser.add_argument("--alg", default="rsa-sha256")
    args = parser.parse_args()

    payload = json.loads(Path(args.input).read_text())
    payload["issued_at"] = payload.get("issued_at") or utcnow_iso()
    payload["signature_alg"] = args.alg
    body = canonical_payload(payload)

    with tempfile.NamedTemporaryFile(delete=False) as data_file, tempfile.NamedTemporaryFile(delete=False) as sig_file:
        data_file.write(body)
        data_file.flush()
        data_path = data_file.name
        sig_path = sig_file.name
    try:
        subprocess.run(
            ["openssl", "dgst", "-sha256", "-sign", args.private_key, "-out", sig_path, data_path],
            check=True,
        )
        signature = base64.b64encode(Path(sig_path).read_bytes()).decode()
    finally:
        Path(data_path).unlink(missing_ok=True)
        Path(sig_path).unlink(missing_ok=True)

    payload["signature"] = signature
    payload["signature_status"] = "signed"
    Path(args.output).write_text(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()