#!/usr/bin/env python3
"""
Print a partner JWT that /connector/amazon/orders accepts — for manual
testing (Postman/curl) without going through onerx-service's auth API.

The connector validates tokens by decrypting the vendorUid claim and
comparing it with RX_API_KEY (no signature check), exactly like
onerx-service issues them. Uses TOKEN_SECRET_KEY / TOKEN_SALT / RX_API_KEY
from the environment when set, else the dev defaults.

Usage:
    python scripts/generate_partner_token.py            # 1h expiry
    python scripts/generate_partner_token.py --hours 8
"""

import argparse
import base64
import hashlib
import json
import os
import time

from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

SECRET_KEY = os.getenv("TOKEN_SECRET_KEY", "aw345-der856-45ty675eu-&8%$#-87967-#@!^7!")
SALT = os.getenv("TOKEN_SALT", "kte45%623@e#{}&&#$%67")
RX_API_KEY = os.getenv("RX_API_KEY", "amazon-6697f65c-8a8f-482e-bd27-7f4e075609d7")


def encrypt_vendor_uid(api_key: str) -> str:
    key = hashlib.pbkdf2_hmac("sha256", SECRET_KEY.encode(), SALT.encode(), 65536, dklen=32)
    padder = padding.PKCS7(128).padder()
    data = padder.update(api_key.encode()) + padder.finalize()
    enc = Cipher(algorithms.AES(key), modes.CBC(bytes(16))).encryptor()
    return base64.b64encode(enc.update(data) + enc.finalize()).decode()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hours", type=float, default=1.0, help="token validity in hours")
    args = parser.parse_args()

    def b64url(obj):
        return base64.urlsafe_b64encode(json.dumps(obj).encode()).rstrip(b"=").decode()

    token = ".".join(
        [
            b64url({"alg": "HS256", "typ": "JWT"}),
            b64url(
                {
                    "exp": int(time.time() + args.hours * 3600),
                    "vendorId": "amazon-manual-test",
                    "vendorUid": encrypt_vendor_uid(RX_API_KEY),
                }
            ),
            "manual-test-signature",
        ]
    )
    print(token)


if __name__ == "__main__":
    main()