"""
Vrismcost — encrypted local vault.

Holds the part of a tax report the chain can never supply: which transfers were
your own wallet moves, what an untraceable acquisition actually cost, the ETH
rates you assigned, your wallet labels and notes. Answer the correction queue
once and it stays answered.

The file lives on your machine and nowhere else. There is no account, no sync,
no server that could be compelled or breached — losing the passphrase and the
recovery code means losing the contents, and that is the trade being made
deliberately.

Design:

  master key      AES-256, generated once from os.urandom
  passphrase wrap scrypt (N=2^17, r=8, p=1) derives a key-encryption key
  recovery wrap   a second, independently generated 32-byte secret wraps the
                  same master key, so a forgotten passphrase is survivable
  records         AES-256-GCM, fresh 96-bit nonce per write
  binding         the vault id is passed as AEAD associated data, so a record
                  lifted from one vault will not decrypt into another
  writes          atomic: temp file, fsync, replace

scrypt at N=2^17 costs ~128MB and roughly a tenth of a second per attempt,
which is what makes a stolen vault expensive to attack offline. That is also
why unlocking is not instant.

Requires `cryptography` — the one dependency in this project. AES-GCM must come
from an audited implementation; composing one out of hashlib would be the kind
of cleverness that loses somebody their records.
"""

from __future__ import annotations

import base64
import json
import os
import secrets
import time
from dataclasses import dataclass
from hashlib import scrypt
from pathlib import Path
from typing import Any

try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
except ImportError:  # pragma: no cover
    raise SystemExit(
        "Vrismcost needs the 'cryptography' package for the encrypted vault.\n"
        "  pip install cryptography"
    )

FORMAT = "vrismcost-vault-1"
SCRYPT_N = 1 << 17
SCRYPT_R = 8
SCRYPT_P = 1
KEY_LEN = 32
NONCE_LEN = 12
RECOVERY_PREFIX = "VRSM"

DEFAULT_PATH = Path.home() / ".vrismcost" / "vault.json"


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


def unb64(text: str) -> bytes:
    return base64.b64decode(text)


class VaultError(Exception):
    pass


class WrongPassphrase(VaultError):
    pass


def _derive(passphrase: str, salt: bytes) -> bytes:
    return scrypt(passphrase.encode("utf-8"), salt=salt,
                  n=SCRYPT_N, r=SCRYPT_R, p=SCRYPT_P, dklen=KEY_LEN,
                  maxmem=256 * 1024 * 1024)


def _aad(vault_id: str, purpose: str) -> bytes:
    """Binds ciphertext to this vault and this role. A wrapped key cannot be
    replayed as a record, and neither can cross vaults."""
    return f"{FORMAT}|{vault_id}|{purpose}".encode()


def format_recovery(secret: bytes) -> str:
    """Grouped base32, easy to write down and read back without ambiguity."""
    body = base64.b32encode(secret).decode().rstrip("=")
    groups = [body[i:i + 5] for i in range(0, len(body), 5)]
    return RECOVERY_PREFIX + "-" + "-".join(groups)


def parse_recovery(code: str) -> bytes:
    raw = code.strip().upper().replace(" ", "").replace("-", "")
    if raw.startswith(RECOVERY_PREFIX):
        raw = raw[len(RECOVERY_PREFIX):]
    pad = "=" * (-len(raw) % 8)
    try:
        secret = base64.b32decode(raw + pad)
    except Exception:
        raise VaultError("That recovery code is not readable.") from None
    if len(secret) != KEY_LEN:
        raise VaultError("That recovery code is the wrong length.")
    return secret


def validate_passphrase(passphrase: str) -> None:
    if len(passphrase) < 10:
        raise VaultError("Use at least 10 characters. A vault backup can be "
                         "attacked offline for as long as someone likes.")


@dataclass
class Vault:
    path: Path
    meta: dict
    key: bytes
    data: dict

    # ---------------------------------------------------------------- create
    @staticmethod
    def create(passphrase: str, path: Path = DEFAULT_PATH) -> tuple["Vault", str]:
        validate_passphrase(passphrase)
        path = Path(path)
        if path.exists():
            raise VaultError(f"A vault already exists at {path}")

        vault_id = secrets.token_hex(16)
        master = os.urandom(KEY_LEN)

        salt = os.urandom(16)
        kek = _derive(passphrase, salt)
        wrapped_nonce = os.urandom(NONCE_LEN)
        wrapped = AESGCM(kek).encrypt(wrapped_nonce, master, _aad(vault_id, "passphrase"))

        recovery_secret = os.urandom(KEY_LEN)
        rec_nonce = os.urandom(NONCE_LEN)
        rec_wrapped = AESGCM(recovery_secret).encrypt(
            rec_nonce, master, _aad(vault_id, "recovery"))

        meta = {
            "format": FORMAT,
            "vaultId": vault_id,
            "created": int(time.time()),
            "revision": 0,
            "kdf": {"name": "scrypt", "n": SCRYPT_N, "r": SCRYPT_R,
                    "p": SCRYPT_P, "salt": b64(salt)},
            "wrappedKey": {"nonce": b64(wrapped_nonce), "ciphertext": b64(wrapped)},
            "recovery": {"nonce": b64(rec_nonce), "ciphertext": b64(rec_wrapped)},
        }
        data = {"wallets": [], "answers": {}, "rates": {}, "settings": {},
                "notes": {}}
        v = Vault(path=path, meta=meta, key=master, data=data)
        v.save()
        return v, format_recovery(recovery_secret)

    # ---------------------------------------------------------------- unlock
    @staticmethod
    def _read_envelope(path: Path) -> dict:
        path = Path(path)
        if not path.exists():
            raise VaultError(f"No vault at {path}. Run `vrismcost vault init` first.")
        env = json.loads(path.read_text(encoding="utf-8"))
        if env.get("format") != FORMAT:
            raise VaultError("That file is not a Vrismcost vault.")
        return env

    @staticmethod
    def unlock(passphrase: str, path: Path = DEFAULT_PATH) -> "Vault":
        env = Vault._read_envelope(path)
        kdf = env["kdf"]
        kek = scrypt(passphrase.encode("utf-8"), salt=unb64(kdf["salt"]),
                     n=kdf["n"], r=kdf["r"], p=kdf["p"], dklen=KEY_LEN,
                     maxmem=256 * 1024 * 1024)
        try:
            master = AESGCM(kek).decrypt(
                unb64(env["wrappedKey"]["nonce"]),
                unb64(env["wrappedKey"]["ciphertext"]),
                _aad(env["vaultId"], "passphrase"))
        except Exception:
            raise WrongPassphrase("That passphrase does not unlock this vault.") from None
        return Vault(path=Path(path), meta=env,
                     key=master, data=Vault._decrypt_body(env, master))

    @staticmethod
    def recover(recovery_code: str, new_passphrase: str,
                path: Path = DEFAULT_PATH) -> tuple["Vault", str]:
        """Recovery rotates the code it consumed. A used code must not stay live."""
        validate_passphrase(new_passphrase)
        env = Vault._read_envelope(path)
        secret = parse_recovery(recovery_code)
        try:
            master = AESGCM(secret).decrypt(
                unb64(env["recovery"]["nonce"]),
                unb64(env["recovery"]["ciphertext"]),
                _aad(env["vaultId"], "recovery"))
        except Exception:
            raise VaultError("That recovery code does not match this vault.") from None

        data = Vault._decrypt_body(env, master)
        salt = os.urandom(16)
        kek = _derive(new_passphrase, salt)
        nonce = os.urandom(NONCE_LEN)
        env["kdf"] = {"name": "scrypt", "n": SCRYPT_N, "r": SCRYPT_R,
                      "p": SCRYPT_P, "salt": b64(salt)}
        env["wrappedKey"] = {
            "nonce": b64(nonce),
            "ciphertext": b64(AESGCM(kek).encrypt(
                nonce, master, _aad(env["vaultId"], "passphrase"))),
        }
        next_secret = os.urandom(KEY_LEN)
        rnonce = os.urandom(NONCE_LEN)
        env["recovery"] = {
            "nonce": b64(rnonce),
            "ciphertext": b64(AESGCM(next_secret).encrypt(
                rnonce, master, _aad(env["vaultId"], "recovery"))),
        }
        v = Vault(path=Path(path), meta=env, key=master, data=data)
        v.save()
        return v, format_recovery(next_secret)

    @staticmethod
    def _decrypt_body(env: dict, master: bytes) -> dict:
        body = env.get("body")
        if not body:
            return {"wallets": [], "answers": {}, "rates": {}, "settings": {}, "notes": {}}
        try:
            plain = AESGCM(master).decrypt(
                unb64(body["nonce"]), unb64(body["ciphertext"]),
                _aad(env["vaultId"], "body"))
        except Exception:
            raise VaultError("The vault contents failed their integrity check. "
                             "The file has been altered or truncated.") from None
        return json.loads(plain.decode("utf-8"))

    # ---------------------------------------------------------------- persist
    def save(self) -> None:
        nonce = os.urandom(NONCE_LEN)
        plain = json.dumps(self.data, separators=(",", ":"), sort_keys=True).encode()
        self.meta["revision"] = int(self.meta.get("revision", 0)) + 1
        self.meta["updated"] = int(time.time())
        self.meta["body"] = {
            "nonce": b64(nonce),
            "ciphertext": b64(AESGCM(self.key).encrypt(
                nonce, plain, _aad(self.meta["vaultId"], "body"))),
        }
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        # Atomic: a crash mid-write must not leave a half vault behind.
        with open(tmp, "w", encoding="utf-8") as fh:
            json.dump(self.meta, fh, indent=2)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, self.path)
        try:
            os.chmod(self.path, 0o600)
        except OSError:
            pass          # best effort; Windows ACLs do not map cleanly

    def lock(self) -> None:
        """Drop the key from memory. Not a guarantee — Python strings and bytes
        are not wipeable — but it stops casual reuse after a task finishes."""
        self.key = b"\x00" * KEY_LEN
        self.data = {}

    # ---------------------------------------------------------------- content
    def add_wallet(self, address: str, label: str = "") -> None:
        address = address.lower()
        for w in self.data["wallets"]:
            if w["address"] == address:
                w["label"] = label or w["label"]
                return
        self.data["wallets"].append({"address": address, "label": label})

    def wallets(self) -> list[str]:
        return [w["address"] for w in self.data["wallets"]]

    def answer(self, tx: str, action: str, value: str | None = None,
               note: str = "", date: str | None = None) -> None:
        """One answer to one queued transaction, stamped as user-supplied so the
        export can always separate it from chain evidence. date, YYYY-MM-DD, is
        when you really acquired what arrived -- a transfer from your own wallet
        keeps the date you first bought it, not the day it moved."""
        self.data["answers"][tx.lower()] = {
            "action": action, "value": value, "note": note, "date": date,
            "at": int(time.time()), "source": "user",
        }

    def unanswer(self, tx: str) -> None:
        self.data["answers"].pop(tx.lower(), None)

    def answers(self) -> dict[str, Any]:
        return self.data["answers"]

    def set_rate(self, symbol: str, date: str, value: str) -> None:
        self.data["rates"].setdefault(symbol.upper(), {})[date] = value

    def rate(self, symbol: str, date: str) -> str | None:
        return self.data["rates"].get(symbol.upper(), {}).get(date)

    def setting(self, key: str, value: Any = None) -> Any:
        if value is None:
            return self.data["settings"].get(key)
        self.data["settings"][key] = value
        return value

    def summary(self) -> dict:
        return {
            "path": str(self.path),
            "vaultId": self.meta["vaultId"],
            "revision": self.meta["revision"],
            "wallets": len(self.data["wallets"]),
            "answers": len(self.data["answers"]),
            "rates": sum(len(v) for v in self.data["rates"].values()),
            "bytes": self.path.stat().st_size if self.path.exists() else 0,
        }