#!/usr/bin/env python3
"""
manage-passwords.py: Passphrase-protected local password vault using GPG.

Features:
- Encrypts/decrypts the vault using a GPG private key protected by a passphrase.
- All decryption and encryption is done in-memory via pipes (no unencrypted data on disk).
- Commands:
    - init: Set up vault with a GPG recipient.
    - import-firefox: Ingest exported Firefox logins CSV.
    - find: Search for entries by domain, URL, or username.
    - update: Update or add a password for a domain/username.
    - list: Display all saved entries (without passwords).
    - delete: Remove an entry.
"""

import argparse
import csv
import getpass
import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

DEFAULT_VAULT_PATH = Path.home() / ".local" / "share" / "password-vault" / "vault.json.gpg"
DEFAULT_CONFIG_PATH = Path.home() / ".local" / "share" / "password-vault" / "config.json"


def get_default_gpg_id():
    """Attempt to find the default secret key id or read from config."""
    if DEFAULT_CONFIG_PATH.exists():
        try:
            with open(DEFAULT_CONFIG_PATH, "r", encoding="utf-8") as f:
                data = json.load(f)
                return data.get("gpg_id")
        except Exception:
            pass
    return None


def save_config(gpg_id, vault_path):
    DEFAULT_CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(DEFAULT_CONFIG_PATH, "w", encoding="utf-8") as f:
        json.dump({"gpg_id": gpg_id, "vault_path": str(vault_path)}, f, indent=2)


def run_gpg_decrypt(vault_path):
    """Decrypt the vault in-memory using gpg."""
    if not vault_path.exists():
        return {}
    cmd = ["gpg", "--quiet", "--batch", "--decrypt", str(vault_path)]
    try:
        proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return json.loads(proc.stdout.decode("utf-8"))
    except subprocess.CalledProcessError:
        # Retry without --batch in case pinentry needs a terminal tty
        cmd_interactive = ["gpg", "--quiet", "--decrypt", str(vault_path)]
        proc = subprocess.run(cmd_interactive, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return json.loads(proc.stdout.decode("utf-8"))


def run_gpg_encrypt(data, vault_path, gpg_id):
    """Encrypt JSON data directly to vault_path atomically."""
    vault_path.parent.mkdir(parents=True, exist_ok=True)
    payload = json.dumps(data, indent=2).encode("utf-8")

    # Atomic write to temporary encrypted file in the same directory
    temp_file = vault_path.with_suffix(".tmp")
    cmd = ["gpg", "--yes", "--quiet", "--batch", "--recipient", gpg_id, "--encrypt", "--output", str(temp_file)]
    try:
        proc = subprocess.run(cmd, input=payload, stderr=subprocess.PIPE, check=True)
        temp_file.replace(vault_path)
    except subprocess.CalledProcessError as e:
        if temp_file.exists():
            temp_file.unlink()
        print(f"Encryption failed: {e.stderr.decode('utf-8')}", file=sys.stderr)
        sys.exit(1)


def cmd_init(args):
    gpg_id = args.recipient
    vault_path = Path(args.vault).expanduser().resolve()

    if not gpg_id:
        # Prompt for GPG key / email
        print("Please enter the GPG key ID or email associated with your passphrase-protected private key.")
        gpg_id = input("GPG Key ID / Email: ").strip()
        if not gpg_id:
            print("Error: GPG key ID cannot be empty.", file=sys.stderr)
            sys.exit(1)

    save_config(gpg_id, vault_path)
    if not vault_path.exists():
        run_gpg_encrypt({}, vault_path, gpg_id)
        print(f"Initialized empty vault at: {vault_path}")
    else:
        print(f"Vault already exists at: {vault_path}. Config updated to use key: {gpg_id}")


def cmd_import_firefox(args):
    csv_path = Path(args.csv_file).expanduser().resolve()
    if not csv_path.exists():
        print(f"Error: Firefox CSV file '{csv_path}' not found.", file=sys.stderr)
        sys.exit(1)

    gpg_id = args.recipient or get_default_gpg_id()
    if not gpg_id:
        print("Error: No GPG recipient configured. Run 'init' first or pass --recipient <email/key_id>.", file=sys.stderr)
        sys.exit(1)

    vault_path = Path(args.vault).expanduser().resolve()
    vault = run_gpg_decrypt(vault_path) if vault_path.exists() else {}

    count = 0
    with open(csv_path, "r", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        for row in reader:
            # Firefox exports headers: url, username, password, httpRealm, formActionOrigin, guid, timeCreated, timeLastUsed, timePasswordChanged
            url = row.get("url") or row.get("hostname") or "unknown"
            username = row.get("username") or ""
            password = row.get("password") or ""

            if not url or not password:
                continue

            if url not in vault:
                vault[url] = {}
            vault[url][username] = {
                "password": password,
                "notes": row.get("httpRealm", "")
            }
            count += 1

    run_gpg_encrypt(vault, vault_path, gpg_id)
    print(f"Successfully imported {count} credentials into {vault_path}.")

    if args.shred:
        try:
            subprocess.run(["shred", "-u", str(csv_path)], check=True)
            print(f"Securely erased {csv_path}.")
        except Exception:
            csv_path.unlink(missing_ok=True)
            print(f"Deleted {csv_path}.")


def cmd_find(args):
    vault_path = Path(args.vault).expanduser().resolve()
    if not vault_path.exists():
        print(f"Vault does not exist at {vault_path}. Run 'init' first.", file=sys.stderr)
        sys.exit(1)

    vault = run_gpg_decrypt(vault_path)
    query = args.query.lower()
    matches = []

    for site, accounts in vault.items():
        for username, data in accounts.items():
            if query in site.lower() or query in username.lower():
                matches.append((site, username, data["password"], data.get("notes", "")))

    if not matches:
        print(f"No credentials found matching '{args.query}'.")
        return

    print(f"Found {len(matches)} match(es):\n" + "-" * 50)
    for site, username, password, notes in matches:
        print(f"Site/URL : {site}")
        print(f"Username : {username}")
        if args.show_password:
            print(f"Password : {password}")
        else:
            masked = password[:2] + "*" * max(4, len(password) - 4) + password[-2:] if len(password) > 4 else "****"
            print(f"Password : {masked} (use -s / --show-password to reveal)")
        if notes:
            print(f"Notes    : {notes}")
        print("-" * 50)


def cmd_update(args):
    vault_path = Path(args.vault).expanduser().resolve()
    gpg_id = args.recipient or get_default_gpg_id()

    if not gpg_id:
        print("Error: No GPG recipient configured. Run 'init' first or pass --recipient <email/key_id>.", file=sys.stderr)
        sys.exit(1)

    vault = run_gpg_decrypt(vault_path) if vault_path.exists() else {}

    site = args.site.strip()
    username = args.username.strip() if args.username else None

    # Check if site already exists
    matched_site = None
    for s in vault:
        if s.lower() == site.lower():
            matched_site = s
            break

    if not matched_site:
        matched_site = site
        vault[matched_site] = {}

    if not username:
        existing_users = list(vault[matched_site].keys())
        if len(existing_users) == 1:
            username = existing_users[0]
            print(f"Updating existing user '{username}' for site '{matched_site}'.")
        elif len(existing_users) > 1:
            print(f"Multiple usernames found for '{matched_site}': {', '.join(existing_users)}")
            username = input("Enter username to update: ").strip()
        else:
            username = input("Enter username: ").strip()

    if not username:
        print("Error: Username cannot be empty.", file=sys.stderr)
        sys.exit(1)

    password = args.password
    if not password:
        password = getpass.getpass(f"Enter new password for {username}@{matched_site}: ")
        confirm = getpass.getpass("Confirm password: ")
        if password != confirm:
            print("Error: Passwords do not match.", file=sys.stderr)
            sys.exit(1)

    if matched_site not in vault:
        vault[matched_site] = {}

    existing_notes = vault[matched_site].get(username, {}).get("notes", "")
    vault[matched_site][username] = {
        "password": password,
        "notes": existing_notes
    }

    run_gpg_encrypt(vault, vault_path, gpg_id)
    print(f"Password for '{username}@{matched_site}' updated successfully.")


def cmd_list(args):
    vault_path = Path(args.vault).expanduser().resolve()
    if not vault_path.exists():
        print("Vault is empty or not created yet.")
        return

    vault = run_gpg_decrypt(vault_path)
    if not vault:
        print("Vault is empty.")
        return

    print(f"Vault contains {sum(len(v) for v in vault.values())} logins across {len(vault)} sites:")
    for site, accounts in sorted(vault.items()):
        for user in accounts:
            print(f"  • {site} (user: {user})")


def cmd_delete(args):
    vault_path = Path(args.vault).expanduser().resolve()
    gpg_id = args.recipient or get_default_gpg_id()
    if not gpg_id:
        print("Error: No GPG recipient configured.", file=sys.stderr)
        sys.exit(1)

    vault = run_gpg_decrypt(vault_path)
    site = args.site
    username = args.username

    matched_site = None
    for s in vault:
        if s.lower() == site.lower():
            matched_site = s
            break

    if not matched_site:
        print(f"Site '{site}' not found in vault.")
        return

    if username:
        if username in vault[matched_site]:
            del vault[matched_site][username]
            if not vault[matched_site]:
                del vault[matched_site]
            run_gpg_encrypt(vault, vault_path, gpg_id)
            print(f"Deleted '{username}' under '{matched_site}'.")
        else:
            print(f"Username '{username}' not found under '{matched_site}'.")
    else:
        del vault[matched_site]
        run_gpg_encrypt(vault, vault_path, gpg_id)
        print(f"Deleted all logins for '{matched_site}'.")


def main():
    default_vault = str(DEFAULT_VAULT_PATH)
    default_gpg = get_default_gpg_id()

    parent_parser = argparse.ArgumentParser(add_help=False)
    parent_parser.add_argument("--vault", default=default_vault, help=f"Path to encrypted vault file (default: {default_vault})")
    parent_parser.add_argument("--recipient", "-r", default=default_gpg, help="GPG key ID / Email used for encryption")

    parser = argparse.ArgumentParser(
        description="GPG passphrase-protected local password vault manager.",
        parents=[parent_parser]
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # init
    p_init = subparsers.add_parser("init", parents=[parent_parser], help="Initialize or configure the vault")
    p_init.set_defaults(func=cmd_init)

    # import-firefox
    p_import = subparsers.add_parser("import-firefox", parents=[parent_parser], help="Import logins from Firefox CSV export")
    p_import.add_argument("csv_file", help="Path to exported Firefox CSV file")
    p_import.add_argument("--shred", action="store_true", help="Securely delete CSV file after import")
    p_import.set_defaults(func=cmd_import_firefox)

    # find
    p_find = subparsers.add_parser("find", parents=[parent_parser], help="Find passwords matching site or username")
    p_find.add_argument("query", help="Search keyword (domain, URL, or username)")
    p_find.add_argument("-s", "--show-password", action="store_true", help="Display password in plain text")
    p_find.set_defaults(func=cmd_find)

    # update
    p_update = subparsers.add_parser("update", parents=[parent_parser], help="Update or add a password")
    p_update.add_argument("site", help="Site / URL (e.g. github.com)")
    p_update.add_argument("-u", "--username", help="Username / Email")
    p_update.add_argument("-p", "--password", help="New password (if omitted, prompts securely)")
    p_update.set_defaults(func=cmd_update)

    # list
    p_list = subparsers.add_parser("list", parents=[parent_parser], help="List all sites and accounts")
    p_list.set_defaults(func=cmd_list)

    # delete
    p_del = subparsers.add_parser("delete", parents=[parent_parser], help="Delete a site or account")
    p_del.add_argument("site", help="Site / URL to delete")
    p_del.add_argument("-u", "--username", help="Specific username to delete (optional)")
    p_del.set_defaults(func=cmd_delete)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()