#!/usr/bin/env python3
"""
age key generator with BIP39 mnemonic backup.

generate  — creates a new X25519 age identity and prints 24 recovery words
recover   — re-creates the exact same age identity from those 24 words
"""

import argparse
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from mnemonic import Mnemonic
import bech32 as bech32lib


ENTROPY_BITS = 256  # → 24 BIP39 words


def encode_identity(private_bytes: bytes) -> str:
    data = bech32lib.convertbits(list(private_bytes), 8, 5)
    return bech32lib.bech32_encode("age-secret-key-", data).upper()


def encode_recipient(public_bytes: bytes) -> str:
    data = bech32lib.convertbits(list(public_bytes), 8, 5)
    return bech32lib.bech32_encode("age", data)


def derive_age_keys(entropy: bytes) -> tuple[str, str]:
    private_key = X25519PrivateKey.from_private_bytes(entropy)
    identity = encode_identity(private_key.private_bytes_raw())
    recipient = encode_recipient(private_key.public_key().public_bytes_raw())
    return identity, recipient


def save_identity_file(identity: str, recipient: str, path: Path) -> None:
    created = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S+00:00")
    content = f"# created: {created}\n# public key: {recipient}\n{identity}\n"
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    path.write_text(content)
    path.chmod(0o600)


def format_mnemonic(words: str) -> str:
    word_list = words.split()
    lines = []
    for i, word in enumerate(word_list, 1):
        lines.append(f"  {i:2d}. {word}")
        if i % 6 == 0 and i < len(word_list):
            lines.append("")
    return "\n".join(lines)


def resolve_output_path(requested: Path) -> Path:
    path = requested
    while path.exists():
        print(f"File already exists: {path}")
        print("  [o] Overwrite  [r] Rename  [q] Quit")
        choice = input("Choice: ").strip().lower()
        if choice == "o":
            break
        elif choice == "r":
            new_name = input(f"New filename (in {path.parent}): ").strip()
            if not new_name:
                print("No name entered, try again.")
                continue
            path = path.parent / new_name
        elif choice == "q":
            print("Aborted.")
            sys.exit(0)
        else:
            print("Please enter o, r, or q.")
    return path


def cmd_generate(args: argparse.Namespace) -> None:
    mnemo = Mnemonic("english")
    entropy = os.urandom(ENTROPY_BITS // 8)
    words = mnemo.to_mnemonic(entropy)
    identity, recipient = derive_age_keys(entropy)

    output = resolve_output_path(Path(args.output))
    save_identity_file(identity, recipient, output)

    print(f"\nGenerated X25519 age identity")
    print(f"  Identity file : {output}")
    print(f"  Public key    : {recipient}")

    print(f"\nYour 24-word recovery mnemonic — write these down and store offline:\n")
    print(format_mnemonic(words))
    print("\nThese words reconstruct your exact age identity. Keep them secret.")


def cmd_recover(args: argparse.Namespace) -> None:
    mnemo = Mnemonic("english")

    if args.words:
        raw = " ".join(args.words)
    else:
        print("Enter your 24 recovery words (space-separated), then press Enter:")
        raw = input().strip()

    words = " ".join(raw.lower().split())

    if not mnemo.check(words):
        print("Error: invalid mnemonic — check spelling or word count.", file=sys.stderr)
        sys.exit(1)

    entropy = bytes(mnemo.to_entropy(words))
    identity, recipient = derive_age_keys(entropy)

    output = resolve_output_path(Path(args.output))
    save_identity_file(identity, recipient, output)

    print(f"\nRecovered X25519 age identity")
    print(f"  Identity file : {output}")
    print(f"  Public key    : {recipient}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate or recover age X25519 identity keys using a BIP39 mnemonic seed",
    )
    sub = parser.add_subparsers(dest="command", required=True, metavar="COMMAND")

    shared = argparse.ArgumentParser(add_help=False)
    shared.add_argument(
        "-o", "--output",
        default=str(Path.home() / ".config" / "age" / "key.txt"),
        metavar="PATH",
        help="output path for the identity file (default: ~/.config/age/key.txt)",
    )

    sub.add_parser(
        "generate",
        aliases=["gen"],
        parents=[shared],
        help="generate a new age identity and print its 24-word mnemonic",
    )
    sub.add_parser(
        "recover",
        aliases=["rec"],
        parents=[shared],
        help="recover an age identity from a 24-word mnemonic",
    ).add_argument(
        "words",
        nargs="*",
        metavar="WORD",
        help="the 24 mnemonic words (omit to be prompted)",
    )

    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.command in ("generate", "gen"):
        cmd_generate(args)
    elif args.command in ("recover", "rec"):
        cmd_recover(args)


if __name__ == "__main__":
    main()