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

generate  — creates a new Ed25519 SSH key pair and prints 24 recovery words
recover   — re-creates the exact same key pair from those 24 words
"""

import argparse
import getpass
import os
import sys
from pathlib import Path

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from mnemonic import Mnemonic


ENTROPY_BITS = 256  # → 24 BIP39 words


def save_keypair(
    private_key: Ed25519PrivateKey,
    path: Path,
    comment: str,
    passphrase: bytes | None,
) -> tuple[Path, Path]:
    encryption = (
        serialization.BestAvailableEncryption(passphrase)
        if passphrase
        else serialization.NoEncryption()
    )

    private_bytes = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.OpenSSH,
        encryption_algorithm=encryption,
    )

    public_bytes = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.OpenSSH,
        format=serialization.PublicFormat.OpenSSH,
    )
    if comment:
        public_bytes = public_bytes + b" " + comment.encode()

    pub_path = path.with_suffix("").parent / (path.name + ".pub")

    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    path.write_bytes(private_bytes)
    path.chmod(0o600)
    pub_path.write_bytes(public_bytes + b"\n")
    pub_path.chmod(0o644)

    return path, pub_path


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 ask_passphrase(confirm: bool = True) -> bytes | None:
    passphrase = getpass.getpass("Key passphrase (leave blank for none): ")
    if not passphrase:
        return None
    if confirm:
        again = getpass.getpass("Confirm passphrase: ")
        if passphrase != again:
            print("Error: passphrases do not match.", file=sys.stderr)
            sys.exit(1)
    return passphrase.encode()


def cmd_generate(args: argparse.Namespace) -> None:
    mnemo = Mnemonic("english")
    entropy = os.urandom(ENTROPY_BITS // 8)
    words = mnemo.to_mnemonic(entropy)
    private_key = Ed25519PrivateKey.from_private_bytes(entropy)

    passphrase = ask_passphrase(confirm=True) if args.passphrase else None

    output = resolve_output_path(Path(args.output))
    key_path, pub_path = save_keypair(private_key, output, args.comment, passphrase)

    print(f"\nGenerated Ed25519 SSH key pair")
    print(f"  Private key : {key_path}")
    print(f"  Public key  : {pub_path}")
    if passphrase:
        print("  Passphrase  : set")

    print(f"\nYour 24-word recovery mnemonic — write these down and store offline:\n")
    print(format_mnemonic(words))
    print("\nThese words reconstruct your exact private key. 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))
    private_key = Ed25519PrivateKey.from_private_bytes(entropy)

    passphrase = ask_passphrase(confirm=True) if args.passphrase else None

    output = resolve_output_path(Path(args.output))
    key_path, pub_path = save_keypair(private_key, output, args.comment, passphrase)

    print(f"\nRecovered Ed25519 SSH key pair")
    print(f"  Private key : {key_path}")
    print(f"  Public key  : {pub_path}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate or recover Ed25519 SSH 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() / ".ssh" / "id_ed25519"),
        metavar="PATH",
        help="output path for the private key (default: ~/.ssh/id_ed25519)",
    )
    shared.add_argument(
        "-C", "--comment",
        default="",
        metavar="TEXT",
        help="comment embedded in the public key (e.g. your email)",
    )
    shared.add_argument(
        "-p", "--passphrase",
        action="store_true",
        help="encrypt the private key file with a passphrase",
    )

    sub.add_parser(
        "generate",
        aliases=["gen"],
        parents=[shared],
        help="generate a new SSH key pair and print its 24-word mnemonic",
    )
    sub.add_parser(
        "recover",
        aliases=["rec"],
        parents=[shared],
        help="recover an SSH key pair from a 24-word mnemonic + timestamp",
    ).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()