#!/usr/bin/env python3
"""
A friendly CLI for NIP-49 encrypt/decrypt and key conversions.
"""

import os
import sys
import unicodedata
from getpass import getpass

try:
    from bech32 import bech32_decode, bech32_encode, convertbits
except ImportError:
    print("⚠️  Missing required package 'bech32'.\n"
          "Run: pip install bech32\n", file=sys.stderr)
    sys.exit(1)

try:
    from ecdsa import SigningKey, SECP256k1
except ImportError:
    print("⚠️  Missing required package 'ecdsa'.\n"
          "Run: pip install ecdsa\n", file=sys.stderr)
    sys.exit(1)

try:
    from nacl.bindings import crypto_aead_xchacha20poly1305_ietf_encrypt, crypto_aead_xchacha20poly1305_ietf_decrypt
    have_xchacha = True
except ImportError:
    have_xchacha = False

try:
    from cryptography.hazmat.primitives.ciphers.aead import XChaCha20Poly1305
    have_crypto_x = True
except ImportError:
    have_crypto_x = False

if not (have_xchacha or have_crypto_x):
    print("❌ No XChaCha20-Poly1305 implementation available.\n"
          "Install either 'pynacl' or 'cryptography' packages.", file=sys.stderr)
    sys.exit(1)

from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

HRP_NCRYPTSEC = "ncryptsec"
HRP_NSEC = "nsec"
HRP_NPUB = "npub"


def clear_terminal():
    os.system("clear" if os.name != "nt" else "cls")


def pause(prompt="[Enter] to continue"):
    input(f"\n{prompt}")


def _normalize_password(password: str) -> str:
    return unicodedata.normalize('NFKC', password)


def _derive_key(password: str, salt: bytes, logn: int) -> bytes:
    pw_norm = _normalize_password(password).encode('utf-8')
    N = 2 ** logn
    kdf = Scrypt(salt=salt, length=32, n=N, r=8, p=1)
    return kdf.derive(pw_norm)


def _bech32_decode_to_bytes(b32: str, expected_hrp: str) -> bytes:
    hrp, data = bech32_decode(b32)
    if hrp is None or data is None:
        raise ValueError("Invalid Bech32 encoding")
    if hrp != expected_hrp:
        raise ValueError(f"Unexpected bech32 prefix: '{hrp}' (expected '{expected_hrp}')")
    raw = convertbits(data, 5, 8, False)
    if raw is None:
        raise ValueError("Bech32 convertbits failed (5→8)")
    return bytes(raw)


def _bech32_encode_from_bytes(hrp: str, data: bytes) -> str:
    words = convertbits(list(data), 8, 5, True)
    if words is None:
        raise ValueError("Bech32 convertbits failed (8→5)")
    return bech32_encode(hrp, words)


def encrypt_nip49(sec: bytes, password: str, logn: int = 16, ksb: int = 0x02) -> str:
    if len(sec) != 32:
        raise ValueError("Private key must be 32 bytes")
    salt = os.urandom(16)
    key = _derive_key(password, salt, logn)
    nonce = os.urandom(24)
    aad = bytes([ksb])
    if have_xchacha:
        ct_and_tag = crypto_aead_xchacha20poly1305_ietf_encrypt(sec, aad, nonce, key)
    else:
        cipher = XChaCha20Poly1305(key)
        ct_and_tag = cipher.encrypt(nonce, sec, aad)
    version_byte = bytes([0x02])
    logn_byte = bytes([logn])
    payload = version_byte + logn_byte + salt + nonce + aad + ct_and_tag
    return _bech32_encode_from_bytes(HRP_NCRYPTSEC, payload)


def decrypt_nip49(ncryptsec: str, password: str) -> bytes:
    b = _bech32_decode_to_bytes(ncryptsec, HRP_NCRYPTSEC)
    if len(b) < (1 + 1 + 16 + 24 + 1 + 16):
        raise ValueError("Invalid payload length")
    version = b[0]
    logn = b[1]
    if version != 0x02:
        raise ValueError(f"Invalid version {version}, expected 0x02")
    salt = b[2:18]
    nonce = b[18:42]
    ksb = b[42]
    aad = bytes([ksb])
    ciphertext = b[43:]
    key = _derive_key(password, salt, logn)
    if have_xchacha:
        plaintext = crypto_aead_xchacha20poly1305_ietf_decrypt(ciphertext, aad, nonce, key)
    else:
        cipher = XChaCha20Poly1305(key)
        plaintext = cipher.decrypt(nonce, ciphertext, aad)
    if len(plaintext) != 32:
        raise ValueError("Decrypted plaintext length invalid")
    return plaintext


def private_bytes_to_nsec(sec: bytes) -> str:
    return _bech32_encode_from_bytes(HRP_NSEC, sec)


def nsec_to_private_bytes(nsec: str) -> bytes:
    return _bech32_decode_to_bytes(nsec, HRP_NSEC)


def private_key_to_public_key_hex(privkey_bytes: bytes) -> str:
    sk = SigningKey.from_string(privkey_bytes, curve=SECP256k1)
    vk = sk.get_verifying_key()
    return vk.to_string("compressed").hex()


def public_key_to_npub(pubkey_bytes: bytes) -> str:
    return _bech32_encode_from_bytes(HRP_NPUB, pubkey_bytes)


def welcome():
    clear_terminal()
    print("=" * 40)
    print("🔐 NIP-49 Encrypt/Decrypt & Key Conversion CLI")
    print("=" * 40)
    print()


def prompt_mode():
    print("Choose an action:")
    print("  (d) Decrypt an encrypted key with password")
    print("  (e) Encrypt a private key with password")
    print("  (h) Convert private key hex to nsec")
    print("  (p) Convert public key hex to npub")
    print("  (q) Quit")
    return input("Your choice? ").strip().lower()


def main():
    welcome()

    while True:
        mode = prompt_mode()

        if mode == 'q':
            print("\nGoodbye!\n")
            break

        elif mode == 'd':
            enc = input("\nEnter ncryptsec: ").strip()
            pwd = getpass("Password: ")
            try:
                decrypted = decrypt_nip49(enc, pwd)
                print("\nDecrypted raw private key (hex):")
                print(decrypted.hex())
                print("\nBech32 nsec:")
                print(private_bytes_to_nsec(decrypted))
            except Exception as e:
                print(f"Decryption failed: {e}")

            pause()
            clear_terminal()

        elif mode == 'e':
            inp = input("\nEnter your nsec or raw private key hex: ").strip()
            if inp.startswith(HRP_NSEC + "1"):
                try:
                    sec = nsec_to_private_bytes(inp)
                except Exception as e:
                    print(f"Invalid nsec input: {e}")
                    pause()
                    clear_terminal()
                    continue
            else:
                try:
                    sec = bytes.fromhex(inp)
                except Exception:
                    print("Invalid hex input.")
                    pause()
                    clear_terminal()
                    continue
            if len(sec) != 32:
                print("Private key must be exactly 32 bytes.")
                pause()
                clear_terminal()
                continue
            pwd = getpass("Password: ")
            ksb_str = input("Key security byte (0/1/2) [default 2]: ").strip()
            ksb = 0x02
            if ksb_str:
                try:
                    ksb = int(ksb_str)
                except ValueError:
                    print("Invalid input for key security byte; using default 2.")
            logn_str = input("logn [default 16]: ").strip()
            logn = 16
            if logn_str:
                try:
                    logn = int(logn_str)
                except ValueError:
                    print("Invalid logn input; using default 16.")
            try:
                encrypted = encrypt_nip49(sec, pwd, logn=logn, ksb=ksb)
                print("\nEncrypted ncryptsec:")
                print(encrypted)
            except Exception as e:
                print(f"Encryption failed: {e}")

            pause()
            clear_terminal()

        elif mode == 'h':
            hex_key = input("\nEnter private key hex (64 hex chars): ").strip()
            try:
                sec = bytes.fromhex(hex_key)
            except Exception:
                print("Invalid hex input.")
                pause()
                clear_terminal()
                continue
            if len(sec) != 32:
                print("Private key must be exactly 32 bytes.")
                pause()
                clear_terminal()
                continue
            nsec = private_bytes_to_nsec(sec)
            print("\nBech32 nsec:")
            print(nsec)

            pause()
            clear_terminal()

        elif mode == 'p':
            pub_hex = input("\nEnter public key hex (compressed 33 bytes / 66 hex chars or uncompressed 65 bytes / 130 hex chars): ").strip()
            try:
                pub_bytes = bytes.fromhex(pub_hex)
            except Exception:
                print("Invalid hex input.")
                pause()
                clear_terminal()
                continue
            if len(pub_bytes) not in (33, 65):
                print("Public key must be 33 bytes (compressed) or 65 bytes (uncompressed).")
                pause()
                clear_terminal()
                continue
            try:
                npub = public_key_to_npub(pub_bytes)
                print("\nBech32 npub:")
                print(npub)
            except Exception as e:
                print(f"Failed to convert public key hex to npub: {e}")

            pause()
            clear_terminal()

        else:
            print("Invalid choice. Please select 'd', 'e', 'h', 'p', or 'q'.")
            pause()
            clear_terminal()


if __name__ == "__main__":
    main()