#!/usr/bin/env python3
"""Local Password Vault

A small educational password manager that:
- Generates passwords with Python's `secrets` module.
- Encrypts the entire vault with AES-256-GCM.
- Derives the encryption key from a master password with scrypt.
- Stores only encrypted data in the user's home folder.

This is a learning project, not a replacement for an audited password manager.
"""

from __future__ import annotations

import base64
import json
import os
import secrets
import string
import tkinter as tk
from datetime import datetime, timezone
from pathlib import Path
from tkinter import messagebox, ttk
from typing import Any

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

APP_NAME = "Local Password Vault"
APP_DIR = Path.home() / ".local_password_vault"
VAULT_FILE = APP_DIR / "vault.json"
ASSOCIATED_DATA = b"local-password-vault-v1"

# OWASP's current minimum scrypt configuration when Argon2id is unavailable.
SCRYPT_N = 2**17
SCRYPT_R = 8
SCRYPT_P = 1

PASSWORD_ALPHABETS = {
    "lowercase": string.ascii_lowercase,
    "uppercase": string.ascii_uppercase,
    "numbers": string.digits,
    "symbols": "!@#$%^&*()-_=+[]{};:,.?",
}


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def b64encode(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def b64decode(data: str) -> bytes:
    return base64.b64decode(data.encode("ascii"), validate=True)


def derive_key(
    master_password: str,
    salt: bytes,
    *,
    n: int = SCRYPT_N,
    r: int = SCRYPT_R,
    p: int = SCRYPT_P,
) -> bytes:
    """Derive a 256-bit encryption key from the master password."""
    kdf = Scrypt(salt=salt, length=32, n=n, r=r, p=p)
    return kdf.derive(master_password.encode("utf-8"))


def generate_password(
    length: int = 20,
    *,
    lowercase: bool = True,
    uppercase: bool = True,
    numbers: bool = True,
    symbols: bool = True,
) -> str:
    """Generate a password containing at least one selected character type."""
    selected: list[str] = []
    if lowercase:
        selected.append(PASSWORD_ALPHABETS["lowercase"])
    if uppercase:
        selected.append(PASSWORD_ALPHABETS["uppercase"])
    if numbers:
        selected.append(PASSWORD_ALPHABETS["numbers"])
    if symbols:
        selected.append(PASSWORD_ALPHABETS["symbols"])

    if not selected:
        raise ValueError("Select at least one character type.")
    if length < len(selected):
        raise ValueError(f"Length must be at least {len(selected)}.")

    # Guarantee one character from every selected group.
    characters = [secrets.choice(group) for group in selected]
    combined = "".join(selected)
    characters.extend(secrets.choice(combined) for _ in range(length - len(characters)))

    # Shuffle securely with SystemRandom.
    secrets.SystemRandom().shuffle(characters)
    return "".join(characters)


def encrypted_payload(key: bytes, salt: bytes, entries: list[dict[str, Any]]) -> dict[str, Any]:
    """Return the encrypted vault container ready to save as JSON."""
    nonce = os.urandom(12)
    plaintext = json.dumps(
        {"entries": entries},
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")
    ciphertext = AESGCM(key).encrypt(nonce, plaintext, ASSOCIATED_DATA)

    return {
        "version": 1,
        "kdf": {
            "name": "scrypt",
            "n": SCRYPT_N,
            "r": SCRYPT_R,
            "p": SCRYPT_P,
            "salt": b64encode(salt),
        },
        "cipher": {
            "name": "AES-256-GCM",
            "nonce": b64encode(nonce),
            "ciphertext": b64encode(ciphertext),
        },
    }


def save_vault(path: Path, key: bytes, salt: bytes, entries: list[dict[str, Any]]) -> None:
    """Atomically write the encrypted vault to disk."""
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = encrypted_payload(key, salt, entries)
    temp_path = path.with_suffix(".tmp")

    with temp_path.open("w", encoding="utf-8") as file:
        json.dump(payload, file, indent=2)
        file.flush()
        os.fsync(file.fileno())

    os.replace(temp_path, path)
    try:
        path.chmod(0o600)
    except OSError:
        # Windows does not use Unix permission bits in the same way.
        pass


def create_vault(path: Path, master_password: str) -> tuple[bytes, bytes, list[dict[str, Any]]]:
    salt = os.urandom(16)
    key = derive_key(master_password, salt)
    entries: list[dict[str, Any]] = []
    save_vault(path, key, salt, entries)
    return key, salt, entries


def unlock_vault(path: Path, master_password: str) -> tuple[bytes, bytes, list[dict[str, Any]]]:
    """Decrypt a vault, raising ValueError for invalid or unsupported data."""
    try:
        with path.open("r", encoding="utf-8") as file:
            payload = json.load(file)

        if payload.get("version") != 1:
            raise ValueError("Unsupported vault version.")

        kdf_data = payload["kdf"]
        cipher_data = payload["cipher"]
        if kdf_data.get("name") != "scrypt" or cipher_data.get("name") != "AES-256-GCM":
            raise ValueError("Unsupported vault encryption settings.")

        salt = b64decode(kdf_data["salt"])
        nonce = b64decode(cipher_data["nonce"])
        ciphertext = b64decode(cipher_data["ciphertext"])

        key = derive_key(
            master_password,
            salt,
            n=int(kdf_data["n"]),
            r=int(kdf_data["r"]),
            p=int(kdf_data["p"]),
        )
        plaintext = AESGCM(key).decrypt(nonce, ciphertext, ASSOCIATED_DATA)
        data = json.loads(plaintext.decode("utf-8"))
        entries = data.get("entries", [])
        if not isinstance(entries, list):
            raise ValueError("The vault contents are invalid.")
        return key, salt, entries
    except InvalidTag as error:
        raise ValueError("Incorrect master password or damaged vault.") from error
    except (OSError, KeyError, TypeError, json.JSONDecodeError, base64.binascii.Error) as error:
        raise ValueError("The vault file could not be read.") from error


class EntryDialog(tk.Toplevel):
    def __init__(
        self,
        parent: tk.Misc,
        *,
        title: str,
        entry: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(parent)
        self.title(title)
        self.resizable(False, False)
        self.transient(parent)
        self.grab_set()
        self.result: dict[str, str] | None = None
        entry = entry or {}

        frame = ttk.Frame(self, padding=16)
        frame.grid(sticky="nsew")

        self.service_var = tk.StringVar(value=str(entry.get("service", "")))
        self.username_var = tk.StringVar(value=str(entry.get("username", "")))
        self.password_var = tk.StringVar(value=str(entry.get("password", "")))
        self.show_var = tk.BooleanVar(value=False)
        self.length_var = tk.IntVar(value=20)

        ttk.Label(frame, text="Website or service").grid(row=0, column=0, sticky="w")
        service_entry = ttk.Entry(frame, textvariable=self.service_var, width=44)
        service_entry.grid(row=1, column=0, columnspan=3, sticky="ew", pady=(3, 10))

        ttk.Label(frame, text="Username or email").grid(row=2, column=0, sticky="w")
        ttk.Entry(frame, textvariable=self.username_var, width=44).grid(
            row=3, column=0, columnspan=3, sticky="ew", pady=(3, 10)
        )

        ttk.Label(frame, text="Password").grid(row=4, column=0, sticky="w")
        self.password_entry = ttk.Entry(frame, textvariable=self.password_var, width=32, show="•")
        self.password_entry.grid(row=5, column=0, sticky="ew", pady=(3, 10))
        ttk.Checkbutton(frame, text="Show", variable=self.show_var, command=self.toggle_password).grid(
            row=5, column=1, padx=(8, 0), pady=(3, 10)
        )

        generator = ttk.LabelFrame(frame, text="Password generator", padding=10)
        generator.grid(row=6, column=0, columnspan=3, sticky="ew", pady=(0, 10))
        ttk.Label(generator, text="Length").grid(row=0, column=0, sticky="w")
        ttk.Spinbox(generator, from_=12, to=64, textvariable=self.length_var, width=6).grid(
            row=0, column=1, padx=(6, 10)
        )
        ttk.Button(generator, text="Generate", command=self.make_password).grid(row=0, column=2)

        ttk.Label(frame, text="Notes (optional)").grid(row=7, column=0, sticky="w")
        self.notes_text = tk.Text(frame, width=44, height=5, wrap="word")
        self.notes_text.grid(row=8, column=0, columnspan=3, sticky="ew", pady=(3, 12))
        self.notes_text.insert("1.0", str(entry.get("notes", "")))

        buttons = ttk.Frame(frame)
        buttons.grid(row=9, column=0, columnspan=3, sticky="e")
        ttk.Button(buttons, text="Cancel", command=self.destroy).pack(side="right")
        ttk.Button(buttons, text="Save", command=self.save).pack(side="right", padx=(0, 8))

        self.bind("<Return>", lambda _event: self.save())
        self.bind("<Escape>", lambda _event: self.destroy())
        service_entry.focus_set()
        self.wait_visibility()
        self.geometry(f"+{parent.winfo_rootx() + 60}+{parent.winfo_rooty() + 60}")

    def toggle_password(self) -> None:
        self.password_entry.configure(show="" if self.show_var.get() else "•")

    def make_password(self) -> None:
        try:
            self.password_var.set(generate_password(int(self.length_var.get())))
        except (ValueError, tk.TclError) as error:
            messagebox.showerror("Password Generator", str(error), parent=self)

    def save(self) -> None:
        service = self.service_var.get().strip()
        password = self.password_var.get()
        if not service:
            messagebox.showwarning("Missing Service", "Enter a website or service name.", parent=self)
            return
        if not password:
            messagebox.showwarning("Missing Password", "Enter or generate a password.", parent=self)
            return

        self.result = {
            "service": service,
            "username": self.username_var.get().strip(),
            "password": password,
            "notes": self.notes_text.get("1.0", "end-1c").strip(),
        }
        self.destroy()


class PasswordVaultApp:
    def __init__(self, root: tk.Tk, vault_path: Path = VAULT_FILE) -> None:
        self.root = root
        self.vault_path = vault_path
        self.key: bytes | None = None
        self.salt: bytes | None = None
        self.entries: list[dict[str, Any]] = []
        self.search_var = tk.StringVar()
        self.status_var = tk.StringVar(value="Locked")
        self.clipboard_job: str | None = None
        self.last_copied_password: str | None = None

        root.title(APP_NAME)
        root.geometry("760x500")
        root.minsize(680, 420)
        root.protocol("WM_DELETE_WINDOW", self.close_app)

        self.container = ttk.Frame(root, padding=18)
        self.container.pack(fill="both", expand=True)
        self.show_login_screen()

    def clear_container(self) -> None:
        for child in self.container.winfo_children():
            child.destroy()

    def show_login_screen(self) -> None:
        self.clear_container()
        self.root.geometry("520x360")

        card = ttk.Frame(self.container, padding=24)
        card.place(relx=0.5, rely=0.48, anchor="center")
        ttk.Label(card, text=APP_NAME, font=("TkDefaultFont", 20, "bold")).grid(
            row=0, column=0, columnspan=2, pady=(0, 8)
        )

        if self.vault_path.exists():
            ttk.Label(card, text="Enter your master password to unlock the encrypted vault.").grid(
                row=1, column=0, columnspan=2, pady=(0, 16)
            )
            password_var = tk.StringVar()
            password_entry = ttk.Entry(card, textvariable=password_var, show="•", width=34)
            password_entry.grid(row=2, column=0, columnspan=2, pady=(0, 12))
            ttk.Button(card, text="Unlock Vault", command=lambda: self.unlock(password_var.get())).grid(
                row=3, column=0, columnspan=2, sticky="ew"
            )
            self.root.bind("<Return>", lambda _event: self.unlock(password_var.get()))
            password_entry.focus_set()
        else:
            ttk.Label(card, text="Create a master password for your new local vault.").grid(
                row=1, column=0, columnspan=2, pady=(0, 14)
            )
            master_var = tk.StringVar()
            confirm_var = tk.StringVar()
            ttk.Label(card, text="Master password").grid(row=2, column=0, sticky="w")
            master_entry = ttk.Entry(card, textvariable=master_var, show="•", width=34)
            master_entry.grid(row=3, column=0, columnspan=2, pady=(3, 10))
            ttk.Label(card, text="Confirm master password").grid(row=4, column=0, sticky="w")
            ttk.Entry(card, textvariable=confirm_var, show="•", width=34).grid(
                row=5, column=0, columnspan=2, pady=(3, 12)
            )
            ttk.Button(
                card,
                text="Create Encrypted Vault",
                command=lambda: self.create(master_var.get(), confirm_var.get()),
            ).grid(row=6, column=0, columnspan=2, sticky="ew")
            ttk.Label(
                card,
                text="There is no password recovery. Store the master password safely.",
                foreground="#8a3b12",
                wraplength=330,
            ).grid(row=7, column=0, columnspan=2, pady=(12, 0))
            self.root.bind("<Return>", lambda _event: self.create(master_var.get(), confirm_var.get()))
            master_entry.focus_set()

    def create(self, master_password: str, confirmation: str) -> None:
        if len(master_password) < 12:
            messagebox.showwarning("Master Password", "Use at least 12 characters.")
            return
        if master_password != confirmation:
            messagebox.showwarning("Master Password", "The passwords do not match.")
            return
        try:
            self.key, self.salt, self.entries = create_vault(self.vault_path, master_password)
        except (OSError, ValueError) as error:
            messagebox.showerror("Create Vault", f"The vault could not be created.\n\n{error}")
            return
        self.show_vault_screen()

    def unlock(self, master_password: str) -> None:
        if not master_password:
            messagebox.showwarning("Unlock Vault", "Enter your master password.")
            return
        try:
            self.key, self.salt, self.entries = unlock_vault(self.vault_path, master_password)
        except ValueError as error:
            messagebox.showerror("Unlock Vault", str(error))
            return
        self.show_vault_screen()

    def show_vault_screen(self) -> None:
        self.root.unbind("<Return>")
        self.clear_container()
        self.root.geometry("760x500")

        header = ttk.Frame(self.container)
        header.pack(fill="x", pady=(0, 12))
        ttk.Label(header, text="Password Vault", font=("TkDefaultFont", 18, "bold")).pack(side="left")
        ttk.Button(header, text="Lock", command=self.lock).pack(side="right")
        ttk.Button(header, text="Add Password", command=self.add_entry).pack(side="right", padx=(0, 8))

        search_frame = ttk.Frame(self.container)
        search_frame.pack(fill="x", pady=(0, 10))
        ttk.Label(search_frame, text="Search:").pack(side="left")
        search_entry = ttk.Entry(search_frame, textvariable=self.search_var)
        search_entry.pack(side="left", fill="x", expand=True, padx=(8, 0))
        self.search_var.trace_add("write", lambda *_args: self.refresh_tree())

        tree_frame = ttk.Frame(self.container)
        tree_frame.pack(fill="both", expand=True)
        self.tree = ttk.Treeview(tree_frame, columns=("service", "username"), show="headings", selectmode="browse")
        self.tree.heading("service", text="Website or Service")
        self.tree.heading("username", text="Username or Email")
        self.tree.column("service", width=280, anchor="w")
        self.tree.column("username", width=330, anchor="w")
        scrollbar = ttk.Scrollbar(tree_frame, orient="vertical", command=self.tree.yview)
        self.tree.configure(yscrollcommand=scrollbar.set)
        self.tree.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        self.tree.bind("<Double-1>", lambda _event: self.edit_entry())

        actions = ttk.Frame(self.container)
        actions.pack(fill="x", pady=(10, 0))
        ttk.Button(actions, text="Edit", command=self.edit_entry).pack(side="left")
        ttk.Button(actions, text="Copy Password", command=self.copy_password).pack(side="left", padx=8)
        ttk.Button(actions, text="Delete", command=self.delete_entry).pack(side="left")
        ttk.Label(actions, textvariable=self.status_var).pack(side="right")

        self.status_var.set(f"Unlocked • {len(self.entries)} saved")
        self.refresh_tree()
        search_entry.focus_set()

    def filtered_entries(self) -> list[dict[str, Any]]:
        query = self.search_var.get().strip().lower()
        entries = sorted(self.entries, key=lambda item: str(item.get("service", "")).lower())
        if not query:
            return entries
        return [
            entry
            for entry in entries
            if query in str(entry.get("service", "")).lower()
            or query in str(entry.get("username", "")).lower()
            or query in str(entry.get("notes", "")).lower()
        ]

    def refresh_tree(self) -> None:
        for item in self.tree.get_children():
            self.tree.delete(item)
        for entry in self.filtered_entries():
            self.tree.insert(
                "",
                "end",
                iid=str(entry["id"]),
                values=(entry.get("service", ""), entry.get("username", "")),
            )
        self.status_var.set(f"Unlocked • {len(self.entries)} saved")

    def selected_entry(self) -> dict[str, Any] | None:
        selection = self.tree.selection()
        if not selection:
            messagebox.showinfo("Password Vault", "Select an entry first.")
            return None
        entry_id = selection[0]
        return next((entry for entry in self.entries if str(entry.get("id")) == entry_id), None)

    def add_entry(self) -> None:
        dialog = EntryDialog(self.root, title="Add Password")
        self.root.wait_window(dialog)
        if not dialog.result:
            return
        now = utc_now()
        new_entry: dict[str, Any] = {
            "id": secrets.token_hex(12),
            **dialog.result,
            "created_at": now,
            "updated_at": now,
        }
        self.entries.append(new_entry)
        self.persist("Password saved")

    def edit_entry(self) -> None:
        entry = self.selected_entry()
        if entry is None:
            return
        dialog = EntryDialog(self.root, title="Edit Password", entry=entry)
        self.root.wait_window(dialog)
        if not dialog.result:
            return
        entry.update(dialog.result)
        entry["updated_at"] = utc_now()
        self.persist("Changes saved")

    def delete_entry(self) -> None:
        entry = self.selected_entry()
        if entry is None:
            return
        service = str(entry.get("service", "this entry"))
        if not messagebox.askyesno("Delete Password", f"Delete the saved password for {service}?"):
            return
        self.entries.remove(entry)
        self.persist("Password deleted")

    def copy_password(self) -> None:
        entry = self.selected_entry()
        if entry is None:
            return
        password = str(entry.get("password", ""))
        self.root.clipboard_clear()
        self.root.clipboard_append(password)
        self.root.update_idletasks()
        self.last_copied_password = password
        self.status_var.set("Password copied • clipboard clears in 30 seconds")

        if self.clipboard_job:
            self.root.after_cancel(self.clipboard_job)
        self.clipboard_job = self.root.after(30_000, self.clear_password_clipboard)

    def clear_password_clipboard(self) -> None:
        self.clipboard_job = None
        try:
            if self.last_copied_password and self.root.clipboard_get() == self.last_copied_password:
                self.root.clipboard_clear()
                self.root.update_idletasks()
                self.status_var.set("Clipboard cleared")
        except tk.TclError:
            pass
        finally:
            self.last_copied_password = None

    def persist(self, status: str) -> None:
        if self.key is None or self.salt is None:
            messagebox.showerror("Password Vault", "The vault is locked.")
            return
        try:
            save_vault(self.vault_path, self.key, self.salt, self.entries)
        except OSError as error:
            messagebox.showerror("Save Vault", f"Your changes could not be saved.\n\n{error}")
            return
        self.refresh_tree()
        self.status_var.set(status)

    def lock(self) -> None:
        self.clear_password_clipboard()
        self.key = None
        self.salt = None
        self.entries = []
        self.search_var.set("")
        self.status_var.set("Locked")
        self.show_login_screen()

    def close_app(self) -> None:
        self.clear_password_clipboard()
        self.key = None
        self.salt = None
        self.entries = []
        self.root.destroy()


def main() -> None:
    root = tk.Tk()
    try:
        root.call("tk", "scaling", 1.15)
    except tk.TclError:
        pass
    PasswordVaultApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()