use std::collections::HashSet;
use std::env;
use std::fmt;
use std::fs::{self, OpenOptions};
use std::io::{self, Read};
use std::path::{Path, PathBuf};

use anyhow::{anyhow, bail, Context, Result};
use argon2::{Algorithm, Argon2, Params, Version};
use base32::Alphabet::Rfc4648;
use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine as _;
use chacha20poly1305::aead::{Aead, KeyInit, Payload};
use chacha20poly1305::{Key, XChaCha20Poly1305, XNonce};
use chrono::{DateTime, Utc};
use ciborium::{de::from_reader, ser::into_writer};
use clap::{Args, Parser, Subcommand, ValueEnum};
use directories::BaseDirs;
use is_terminal::IsTerminal;
use kem::{Decapsulate, Encapsulate};
use ml_kem::{array::Array, Ciphertext, Encoded, EncodedSizeUser, KemCore, MlKem1024};
use rand::{rngs::OsRng, RngCore};
use rpassword::prompt_password;
use serde::{Deserialize, Serialize};
use serde_json::json;
use tempfile::NamedTempFile;
use totp_rs::{Algorithm as TotpAlgorithmLib, Secret as TotpSecret, TOTP};
use typenum::Unsigned;
use uuid::Uuid;
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

const VAULT_VERSION: u32 = 1;
const AAD_DEK: &[u8] = b"black-bag::sealed-dek";
const AAD_DK: &[u8] = b"black-bag::sealed-dk";
const AAD_PAYLOAD: &[u8] = b"black-bag::payload";
const DEFAULT_TIME_COST: u32 = 3;
const DEFAULT_LANES: u32 = 1;

pub fn run() -> Result<()> {
    let cli = Cli::parse();
    match cli.command {
        Command::Init(cmd) => init_vault(cmd),
        Command::Add(cmd) => add_record(cmd),
        Command::List(cmd) => list_records(cmd),
        Command::Get(cmd) => get_record(cmd),
        Command::Rotate(cmd) => rotate_vault(cmd),
        Command::Doctor(cmd) => doctor(cmd),
        Command::Recovery { command } => recovery(command),
        Command::Totp { command } => totp(command),
        Command::Selftest => self_test(),
    }
}

#[derive(Parser)]
#[command(name = "black-bag", version, about = "Ultra-secure zero-trace CLI vault", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Initialize a new vault
    Init(InitCommand),
    /// Add a record to the vault
    Add(AddCommand),
    /// List records (masked summaries)
    List(ListCommand),
    /// Inspect a record by UUID
    Get(GetCommand),
    /// Rewrap the master key with fresh randomness
    Rotate(RotateCommand),
    /// Print health diagnostics
    Doctor(DoctorCommand),
    /// Manage Shamir recovery shares
    Recovery {
        #[command(subcommand)]
        command: RecoveryCommand,
    },
    /// Work with stored TOTP secrets
    Totp {
        #[command(subcommand)]
        command: TotpCommand,
    },
    /// Run embedded self-tests
    Selftest,
}

#[derive(Args)]
struct InitCommand {
    /// Argon2 memory cost in KiB (default: 262144 => 256 MiB)
    #[arg(long, default_value_t = 262_144)]
    mem_kib: u32,
}

#[derive(Args)]
struct ListCommand {
    /// Filter by record kind
    #[arg(long, value_enum)]
    kind: Option<RecordKind>,
    /// Filter by tag (case-insensitive substring match)
    #[arg(long)]
    tag: Option<String>,
    /// Full-text query over metadata fields
    #[arg(long)]
    query: Option<String>,
}

#[derive(Args)]
struct GetCommand {
    /// Record UUID to inspect
    id: Uuid,
    /// Reveal sensitive fields (requires TTY)
    #[arg(long)]
    reveal: bool,
}

#[derive(Args, Default)]
struct RotateCommand {
    /// Optionally override Argon2 memory cost in KiB during rotation
    #[arg(long)]
    mem_kib: Option<u32>,
}

#[derive(Args)]
struct DoctorCommand {
    /// Emit machine-readable JSON instead of human text
    #[arg(long)]
    json: bool,
}

#[derive(Subcommand)]
enum RecoveryCommand {
    /// Split a secret into Shamir shares
    Split(RecoverySplitCommand),
    /// Combine Shamir shares back into the original secret
    Combine(RecoveryCombineCommand),
}

#[derive(Args)]
struct RecoverySplitCommand {
    /// Threshold required to reconstruct the secret
    #[arg(long, default_value_t = 3)]
    threshold: u8,
    /// Total number of shares to produce
    #[arg(long, default_value_t = 5)]
    shares: u8,
}

#[derive(Args)]
struct RecoveryCombineCommand {
    /// Reconstruction threshold (usually matches value used during split)
    #[arg(long)]
    threshold: u8,
    /// Comma-separated list of shares (e.g., 1-<base64>,2-<base64>)
    #[arg(long)]
    shares: String,
}

#[derive(Subcommand)]
enum TotpCommand {
    /// Generate a TOTP code for the specified record
    Code(TotpCodeCommand),
}

#[derive(Args)]
struct TotpCodeCommand {
    /// Record UUID containing the TOTP secret
    id: Uuid,
    /// Unix timestamp (seconds) to use instead of now
    #[arg(long)]
    time: Option<i64>,
}

#[derive(Args)]
struct AddCommand {
    #[command(subcommand)]
    record: AddRecord,
}

#[derive(Subcommand)]
enum AddRecord {
    /// Add a login/password record
    Login(AddLogin),
    /// Add a contact record
    Contact(AddContact),
    /// Add an identity document record
    Id(AddIdentity),
    /// Add a secure note
    Note(AddNote),
    /// Add a bank account record
    Bank(AddBank),
    /// Add a Wi-Fi profile record
    Wifi(AddWifi),
    /// Add an API credential record
    Api(AddApi),
    /// Add a cryptocurrency wallet record
    Wallet(AddWallet),
    /// Add a TOTP secret
    Totp(AddTotp),
    /// Add an SSH key record
    Ssh(AddSsh),
    /// Add a PGP key record
    Pgp(AddPgp),
    /// Add a recovery kit record
    Recovery(AddRecovery),
}

#[derive(Args)]
struct CommonRecordArgs {
    /// Optional title for the record
    #[arg(long)]
    title: Option<String>,
    /// Comma-separated list of tags
    #[arg(long, value_delimiter = ',')]
    tags: Vec<String>,
    /// Optional free-form notes (stored alongside metadata)
    #[arg(long)]
    notes: Option<String>,
}

#[derive(Args)]
struct AddLogin {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    username: Option<String>,
    #[arg(long)]
    url: Option<String>,
}

#[derive(Args)]
struct AddContact {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long, required = true)]
    full_name: String,
    #[arg(long, value_delimiter = ',')]
    emails: Vec<String>,
    #[arg(long, value_delimiter = ',')]
    phones: Vec<String>,
}

#[derive(Args)]
struct AddIdentity {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    id_type: Option<String>,
    #[arg(long)]
    name_on_doc: Option<String>,
    #[arg(long)]
    number: Option<String>,
    #[arg(long)]
    issuing_country: Option<String>,
    #[arg(long)]
    expiry: Option<String>,
}

#[derive(Args)]
struct AddNote {
    #[command(flatten)]
    common: CommonRecordArgs,
}

#[derive(Args)]
struct AddBank {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    institution: Option<String>,
    #[arg(long)]
    account_name: Option<String>,
    #[arg(long)]
    routing_number: Option<String>,
}

#[derive(Args)]
struct AddWifi {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    ssid: Option<String>,
    #[arg(long)]
    security: Option<String>,
    #[arg(long)]
    location: Option<String>,
}

#[derive(Args)]
struct AddApi {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    service: Option<String>,
    #[arg(long)]
    environment: Option<String>,
    #[arg(long)]
    access_key: Option<String>,
    #[arg(long, value_delimiter = ',')]
    scopes: Vec<String>,
}

#[derive(Args)]
struct AddWallet {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    asset: Option<String>,
    #[arg(long)]
    address: Option<String>,
    #[arg(long)]
    network: Option<String>,
}

#[derive(Args)]
struct AddTotp {
    #[command(flatten)]
    common: CommonRecordArgs,
    /// Optional issuer string (display only)
    #[arg(long)]
    issuer: Option<String>,
    /// Optional account/name label (display only)
    #[arg(long)]
    account: Option<String>,
    /// Base32-encoded secret; omit to be prompted securely
    #[arg(long)]
    secret: Option<String>,
    /// Number of digits (6-8)
    #[arg(long, default_value_t = 6)]
    digits: u8,
    /// Seconds per step
    #[arg(long, default_value_t = 30)]
    step: u64,
    /// Allowed skew (number of steps on each side)
    #[arg(long, default_value_t = 1)]
    skew: u8,
    /// Hash algorithm
    #[arg(long, value_enum, default_value_t = TotpAlgorithm::Sha1)]
    algorithm: TotpAlgorithm,
}

#[derive(Args)]
struct AddSsh {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    label: Option<String>,
    #[arg(long)]
    comment: Option<String>,
}

#[derive(Args)]
struct AddPgp {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    label: Option<String>,
    #[arg(long)]
    fingerprint: Option<String>,
}

#[derive(Args)]
struct AddRecovery {
    #[command(flatten)]
    common: CommonRecordArgs,
    #[arg(long)]
    description: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ValueEnum, Hash)]
#[serde(rename_all = "snake_case")]
#[clap(rename_all = "snake_case")]
enum RecordKind {
    Login,
    Contact,
    Id,
    Note,
    Bank,
    Wifi,
    Api,
    Wallet,
    Totp,
    Ssh,
    Pgp,
    Recovery,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ValueEnum)]
#[serde(rename_all = "lowercase")]
enum TotpAlgorithm {
    Sha1,
    Sha256,
    Sha512,
}

impl TotpAlgorithm {
    fn to_lib(self) -> TotpAlgorithmLib {
        match self {
            TotpAlgorithm::Sha1 => TotpAlgorithmLib::SHA1,
            TotpAlgorithm::Sha256 => TotpAlgorithmLib::SHA256,
            TotpAlgorithm::Sha512 => TotpAlgorithmLib::SHA512,
        }
    }
}

impl fmt::Display for TotpAlgorithm {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TotpAlgorithm::Sha1 => f.write_str("sha1"),
            TotpAlgorithm::Sha256 => f.write_str("sha256"),
            TotpAlgorithm::Sha512 => f.write_str("sha512"),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
struct Record {
    id: Uuid,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    title: Option<String>,
    tags: Vec<String>,
    metadata_notes: Option<String>,
    data: RecordData,
}

impl Record {
    fn new(
        data: RecordData,
        title: Option<String>,
        tags: Vec<String>,
        notes: Option<String>,
    ) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4(),
            created_at: now,
            updated_at: now,
            title,
            tags,
            metadata_notes: notes,
            data,
        }
    }

    fn kind(&self) -> RecordKind {
        self.data.kind()
    }

    fn matches_tag(&self, tag: &str) -> bool {
        let tag_lower = tag.to_ascii_lowercase();
        self.tags
            .iter()
            .any(|t| t.to_ascii_lowercase().contains(&tag_lower))
    }

    fn matches_query(&self, needle: &str) -> bool {
        let haystack = [
            self.title.as_deref().unwrap_or_default(),
            self.metadata_notes.as_deref().unwrap_or_default(),
            &self.data.summary_text(),
        ]
        .join("\n")
        .to_ascii_lowercase();
        haystack.contains(&needle.to_ascii_lowercase())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind", rename_all = "snake_case")]
enum RecordData {
    Login {
        username: Option<String>,
        url: Option<String>,
        password: Sensitive,
    },
    Contact {
        full_name: String,
        emails: Vec<String>,
        phones: Vec<String>,
    },
    Id {
        id_type: Option<String>,
        name_on_doc: Option<String>,
        number: Option<String>,
        issuing_country: Option<String>,
        expiry: Option<String>,
        secret: Option<Sensitive>,
    },
    Note {
        body: Sensitive,
    },
    Bank {
        institution: Option<String>,
        account_name: Option<String>,
        routing_number: Option<String>,
        account_number: Sensitive,
    },
    Wifi {
        ssid: Option<String>,
        security: Option<String>,
        location: Option<String>,
        passphrase: Sensitive,
    },
    Api {
        service: Option<String>,
        environment: Option<String>,
        access_key: Option<String>,
        secret_key: Sensitive,
        scopes: Vec<String>,
    },
    Wallet {
        asset: Option<String>,
        address: Option<String>,
        network: Option<String>,
        secret_key: Sensitive,
    },
    Totp {
        issuer: Option<String>,
        account: Option<String>,
        secret: Sensitive,
        digits: u8,
        step: u64,
        skew: u8,
        algorithm: TotpAlgorithm,
    },
    Ssh {
        label: Option<String>,
        private_key: Sensitive,
        comment: Option<String>,
    },
    Pgp {
        label: Option<String>,
        fingerprint: Option<String>,
        armored_private_key: Sensitive,
    },
    Recovery {
        description: Option<String>,
        payload: Sensitive,
    },
}

impl RecordData {
    fn kind(&self) -> RecordKind {
        match self {
            RecordData::Login { .. } => RecordKind::Login,
            RecordData::Contact { .. } => RecordKind::Contact,
            RecordData::Id { .. } => RecordKind::Id,
            RecordData::Note { .. } => RecordKind::Note,
            RecordData::Bank { .. } => RecordKind::Bank,
            RecordData::Wifi { .. } => RecordKind::Wifi,
            RecordData::Api { .. } => RecordKind::Api,
            RecordData::Wallet { .. } => RecordKind::Wallet,
            RecordData::Totp { .. } => RecordKind::Totp,
            RecordData::Ssh { .. } => RecordKind::Ssh,
            RecordData::Pgp { .. } => RecordKind::Pgp,
            RecordData::Recovery { .. } => RecordKind::Recovery,
        }
    }

    fn summary_text(&self) -> String {
        match self {
            RecordData::Login { username, url, .. } => format!(
                "user={} url={}",
                username.as_deref().unwrap_or("-"),
                url.as_deref().unwrap_or("-")
            ),
            RecordData::Contact {
                full_name,
                emails,
                phones,
            } => format!(
                "{} | emails={} | phones={}",
                full_name,
                if emails.is_empty() {
                    "-".to_string()
                } else {
                    emails.join(",")
                },
                if phones.is_empty() {
                    "-".to_string()
                } else {
                    phones.join(",")
                }
            ),
            RecordData::Id {
                id_type,
                number,
                expiry,
                ..
            } => format!(
                "type={} number={} expiry={}",
                id_type.as_deref().unwrap_or("-"),
                number.as_deref().unwrap_or("-"),
                expiry.as_deref().unwrap_or("-")
            ),
            RecordData::Note { .. } => "secure note".to_string(),
            RecordData::Bank {
                institution,
                account_name,
                routing_number,
                ..
            } => format!(
                "institution={} account={} routing={}",
                institution.as_deref().unwrap_or("-"),
                account_name.as_deref().unwrap_or("-"),
                routing_number.as_deref().unwrap_or("-")
            ),
            RecordData::Wifi {
                ssid,
                security,
                location,
                ..
            } => format!(
                "ssid={} security={} location={}",
                ssid.as_deref().unwrap_or("-"),
                security.as_deref().unwrap_or("-"),
                location.as_deref().unwrap_or("-")
            ),
            RecordData::Api {
                service,
                environment,
                scopes,
                ..
            } => format!(
                "service={} env={} scopes={}",
                service.as_deref().unwrap_or("-"),
                environment.as_deref().unwrap_or("-"),
                if scopes.is_empty() {
                    "-".to_string()
                } else {
                    scopes.join(",")
                }
            ),
            RecordData::Wallet {
                asset,
                address,
                network,
                ..
            } => format!(
                "asset={} address={} network={}",
                asset.as_deref().unwrap_or("-"),
                address.as_deref().unwrap_or("-"),
                network.as_deref().unwrap_or("-")
            ),
            RecordData::Totp {
                issuer,
                account,
                digits,
                step,
                ..
            } => format!(
                "issuer={} account={} digits={} step={}",
                issuer.as_deref().unwrap_or("-"),
                account.as_deref().unwrap_or("-"),
                digits,
                step
            ),
            RecordData::Ssh { label, comment, .. } => format!(
                "label={} comment={}",
                label.as_deref().unwrap_or("-"),
                comment.as_deref().unwrap_or("-")
            ),
            RecordData::Pgp {
                label, fingerprint, ..
            } => format!(
                "label={} fingerprint={}",
                label.as_deref().unwrap_or("-"),
                fingerprint.as_deref().unwrap_or("-")
            ),
            RecordData::Recovery { description, .. } => {
                format!("description={}", description.as_deref().unwrap_or("-"))
            }
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
struct Sensitive {
    #[serde(with = "serde_bytes")]
    data: Vec<u8>,
}

impl Sensitive {
    fn new_from_utf8(value: &[u8]) -> Self {
        Self {
            data: value.to_vec(),
        }
    }

    fn from_string(value: &str) -> Self {
        Self::new_from_utf8(value.as_bytes())
    }

    fn as_slice(&self) -> &[u8] {
        &self.data
    }

    fn expose_utf8(&self) -> Result<String> {
        Ok(String::from_utf8(self.data.clone())?)
    }
}

impl Drop for Sensitive {
    fn drop(&mut self) {
        self.data.zeroize();
    }
}

impl ZeroizeOnDrop for Sensitive {}

#[derive(Serialize, Deserialize, Clone)]
struct VaultFile {
    version: u32,
    header: VaultHeader,
    payload: AeadBlob,
}

#[derive(Serialize, Deserialize, Clone)]
struct VaultHeader {
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    argon: ArgonState,
    kem_public: Vec<u8>,
    kem_ciphertext: Vec<u8>,
    sealed_decapsulation: AeadBlob,
    sealed_dek: AeadBlob,
}

#[derive(Serialize, Deserialize, Clone)]
struct ArgonState {
    mem_cost_kib: u32,
    time_cost: u32,
    lanes: u32,
    salt: [u8; 32],
}

#[derive(Serialize, Deserialize, Clone)]
struct AeadBlob {
    nonce: [u8; 24],
    ciphertext: Vec<u8>,
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
struct VaultPayload {
    records: Vec<Record>,
    record_counter: u64,
}

struct Vault {
    path: PathBuf,
    file: VaultFile,
    payload: VaultPayload,
    dek: Zeroizing<[u8; 32]>,
}

impl Vault {
    fn init(path: &Path, passphrase: &Zeroizing<String>, mem_kib: u32) -> Result<()> {
        if path.exists() {
            bail!("vault already exists at {}", path.display());
        }

        if mem_kib < 32_768 {
            bail!("mem-kib must be at least 32768 (32 MiB)");
        }

        let mut salt = [0u8; 32];
        OsRng.fill_bytes(&mut salt);

        let argon = ArgonState {
            mem_cost_kib: mem_kib,
            time_cost: DEFAULT_TIME_COST,
            lanes: DEFAULT_LANES,
            salt,
        };

        let kek = derive_kek(passphrase, &argon)?;

        let (dk, ek) = <MlKem1024 as KemCore>::generate(&mut OsRng);
        let (kem_ct, shared_key) = ek
            .encapsulate(&mut OsRng)
            .map_err(|e| anyhow!("ml-kem encapsulate failed: {e:?}"))?;

        let mut dek_bytes = [0u8; 32];
        OsRng.fill_bytes(&mut dek_bytes);
        let dek = Zeroizing::new(dek_bytes);

        let sealed_dek = encrypt_blob(shared_key.as_slice(), dek.as_slice(), AAD_DEK)?;

        let dk_bytes = dk.as_bytes();
        let sealed_decapsulation = encrypt_blob(kek.as_slice(), dk_bytes.as_slice(), AAD_DK)?;

        let payload = VaultPayload {
            records: Vec::new(),
            record_counter: 0,
        };
        let payload_blob = encrypt_payload(dek.as_ref(), &payload)?;

        let header = VaultHeader {
            created_at: Utc::now(),
            updated_at: Utc::now(),
            argon,
            kem_public: ek.as_bytes().to_vec(),
            kem_ciphertext: kem_ct.to_vec(),
            sealed_decapsulation,
            sealed_dek,
        };

        let file = VaultFile {
            version: VAULT_VERSION,
            header,
            payload: payload_blob,
        };

        write_vault(path, &file)
    }

    fn load(path: &Path, passphrase: &Zeroizing<String>) -> Result<Self> {
        if !path.exists() {
            bail!("vault not initialized at {}", path.display());
        }
        let mut file = OpenOptions::new()
            .read(true)
            .open(path)
            .with_context(|| format!("failed to open vault at {}", path.display()))?;

        let mut buf = Zeroizing::new(Vec::new());
        file.read_to_end(&mut buf)?;
        let vault_file: VaultFile = from_reader(buf.as_slice()).context("failed to parse vault")?;

        if vault_file.version != VAULT_VERSION {
            bail!("unsupported vault version {}", vault_file.version);
        }

        let kek = derive_kek(passphrase, &vault_file.header.argon)?;
        let dk_bytes = Zeroizing::new(decrypt_blob(
            kek.as_slice(),
            &vault_file.header.sealed_decapsulation,
            AAD_DK,
        )?);
        type DecapKey = <MlKem1024 as KemCore>::DecapsulationKey;
        let dk_encoded = expect_encoded::<DecapKey>(dk_bytes.as_slice(), "decapsulation key")?;
        let dk = DecapKey::from_bytes(&dk_encoded);

        let kem_ct = expect_ciphertext(&vault_file.header.kem_ciphertext)?;
        let shared = dk
            .decapsulate(&kem_ct)
            .map_err(|_| anyhow!("ml-kem decapsulation failed"))?;

        let dek_bytes = Zeroizing::new(decrypt_blob(
            shared.as_slice(),
            &vault_file.header.sealed_dek,
            AAD_DEK,
        )?);
        if dek_bytes.len() != 32 {
            bail!("invalid dek length");
        }
        let mut dek_array = [0u8; 32];
        dek_array.copy_from_slice(dek_bytes.as_slice());
        let dek = Zeroizing::new(dek_array);

        let payload: VaultPayload = decrypt_payload(dek.as_ref(), &vault_file.payload)?;

        Ok(Self {
            path: path.to_path_buf(),
            file: vault_file,
            payload,
            dek,
        })
    }

    fn save(&mut self, passphrase: &Zeroizing<String>) -> Result<()> {
        self.file.header.updated_at = Utc::now();

        let payload_blob = encrypt_payload(self.dek.as_ref(), &self.payload)?;
        self.file.payload = payload_blob;

        let kek = derive_kek(passphrase, &self.file.header.argon)?;

        let dk_bytes = Zeroizing::new(decrypt_blob(
            kek.as_slice(),
            &self.file.header.sealed_decapsulation,
            AAD_DK,
        )?);
        type DecapKey = <MlKem1024 as KemCore>::DecapsulationKey;
        let dk_encoded = expect_encoded::<DecapKey>(dk_bytes.as_slice(), "decapsulation key")?;
        let dk = DecapKey::from_bytes(&dk_encoded);
        let (kem_ct, shared) = {
            type EncKey = <MlKem1024 as KemCore>::EncapsulationKey;
            let ek_encoded =
                expect_encoded::<EncKey>(&self.file.header.kem_public, "encapsulation key")?;
            let ek = EncKey::from_bytes(&ek_encoded);
            ek.encapsulate(&mut OsRng)
                .map_err(|e| anyhow!("ml-kem encapsulate failed: {e:?}"))?
        };

        let sealed_dek = encrypt_blob(shared.as_slice(), self.dek.as_ref(), AAD_DEK)?;
        let sealed_decapsulation = encrypt_blob(kek.as_slice(), dk.as_bytes().as_slice(), AAD_DK)?;

        self.file.header.kem_ciphertext = kem_ct.to_vec();
        self.file.header.sealed_dek = sealed_dek;
        self.file.header.sealed_decapsulation = sealed_decapsulation;

        write_vault(&self.path, &self.file)
    }

    fn add_record(&mut self, record: Record) {
        self.payload.record_counter = self.payload.record_counter.saturating_add(1);
        self.payload.records.push(record);
    }

    fn list(
        &self,
        kind: Option<RecordKind>,
        tag: Option<&str>,
        query: Option<&str>,
    ) -> Vec<&Record> {
        self.payload
            .records
            .iter()
            .filter(|rec| kind.map(|k| rec.kind() == k).unwrap_or(true))
            .filter(|rec| tag.map(|t| rec.matches_tag(t)).unwrap_or(true))
            .filter(|rec| query.map(|q| rec.matches_query(q)).unwrap_or(true))
            .collect()
    }

    fn get(&mut self, id: Uuid) -> Option<&mut Record> {
        self.payload.records.iter_mut().find(|rec| rec.id == id)
    }

    fn get_ref(&self, id: Uuid) -> Option<&Record> {
        self.payload.records.iter().find(|rec| rec.id == id)
    }

    fn rotate(&mut self, passphrase: &Zeroizing<String>, mem_kib: Option<u32>) -> Result<()> {
        if let Some(mem) = mem_kib {
            if mem < 32_768 {
                bail!("mem-kib must be at least 32768 (32 MiB)");
            }
            self.file.header.argon.mem_cost_kib = mem;
            OsRng.fill_bytes(&mut self.file.header.argon.salt);
        }

        let (dk, ek) = <MlKem1024 as KemCore>::generate(&mut OsRng);
        let (kem_ct, shared_key) = ek
            .encapsulate(&mut OsRng)
            .map_err(|e| anyhow!("ml-kem encapsulate failed: {e:?}"))?;

        let kek = derive_kek(passphrase, &self.file.header.argon)?;
        let sealed_dek = encrypt_blob(shared_key.as_slice(), self.dek.as_ref(), AAD_DEK)?;
        let sealed_decapsulation = encrypt_blob(kek.as_slice(), dk.as_bytes().as_slice(), AAD_DK)?;

        self.file.header.kem_public = ek.as_bytes().to_vec();
        self.file.header.kem_ciphertext = kem_ct.to_vec();
        self.file.header.sealed_dek = sealed_dek;
        self.file.header.sealed_decapsulation = sealed_decapsulation;

        Ok(())
    }

    fn stats(&self) -> VaultStats {
        VaultStats {
            created_at: self.file.header.created_at,
            updated_at: self.file.header.updated_at,
            record_count: self.payload.records.len(),
            argon_mem_kib: self.file.header.argon.mem_cost_kib,
            argon_time_cost: self.file.header.argon.time_cost,
            argon_lanes: self.file.header.argon.lanes,
        }
    }
}

impl Drop for Vault {
    fn drop(&mut self) {
        self.dek.zeroize();
    }
}

struct VaultStats {
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    record_count: usize,
    argon_mem_kib: u32,
    argon_time_cost: u32,
    argon_lanes: u32,
}

fn derive_kek(passphrase: &Zeroizing<String>, argon: &ArgonState) -> Result<Zeroizing<[u8; 32]>> {
    let params = Params::new(argon.mem_cost_kib, argon.time_cost, argon.lanes, Some(32))
        .map_err(|e| anyhow!("invalid Argon2 parameters: {e}"))?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
    let mut output = Zeroizing::new([0u8; 32]);
    argon2
        .hash_password_into(passphrase.as_bytes(), &argon.salt, output.as_mut())
        .map_err(|e| anyhow!("argon2 derivation failed: {e}"))?;
    Ok(output)
}

fn expect_encoded<T>(data: &[u8], label: &str) -> Result<Encoded<T>>
where
    T: EncodedSizeUser,
{
    let expected = <T as EncodedSizeUser>::EncodedSize::USIZE;
    Array::try_from_iter(data.iter().copied()).map_err(|_| {
        anyhow!(
            "invalid {label} length: expected {expected}, got {}",
            data.len()
        )
    })
}

fn expect_ciphertext(data: &[u8]) -> Result<Ciphertext<MlKem1024>> {
    let expected = <MlKem1024 as KemCore>::CiphertextSize::USIZE;
    Array::try_from_iter(data.iter().copied()).map_err(|_| {
        anyhow!(
            "invalid ciphertext length: expected {expected}, got {}",
            data.len()
        )
    })
}

fn encrypt_blob(key: &[u8], plaintext: &[u8], aad: &[u8]) -> Result<AeadBlob> {
    let mut nonce = [0u8; 24];
    OsRng.fill_bytes(&mut nonce);
    let cipher = XChaCha20Poly1305::new(Key::from_slice(key));
    let ciphertext = cipher
        .encrypt(
            XNonce::from_slice(&nonce),
            Payload {
                msg: plaintext,
                aad,
            },
        )
        .map_err(|_| anyhow!("encryption failed"))?;
    Ok(AeadBlob { nonce, ciphertext })
}

fn decrypt_blob(key: &[u8], blob: &AeadBlob, aad: &[u8]) -> Result<Vec<u8>> {
    let cipher = XChaCha20Poly1305::new(Key::from_slice(key));
    let plaintext = cipher
        .decrypt(
            XNonce::from_slice(&blob.nonce),
            Payload {
                msg: &blob.ciphertext,
                aad,
            },
        )
        .map_err(|_| anyhow!("decryption failed"))?;
    Ok(plaintext)
}

fn encrypt_payload(dek: &[u8], payload: &VaultPayload) -> Result<AeadBlob> {
    let mut buf = Vec::new();
    into_writer(payload, &mut buf).context("failed to serialize payload")?;
    encrypt_blob(dek, &buf, AAD_PAYLOAD)
}

fn decrypt_payload(dek: &[u8], blob: &AeadBlob) -> Result<VaultPayload> {
    let plaintext = decrypt_blob(dek, blob, AAD_PAYLOAD)?;
    let payload: VaultPayload =
        from_reader(plaintext.as_slice()).context("failed to parse payload")?;
    Ok(payload)
}

fn vault_path() -> Result<PathBuf> {
    if let Ok(path) = env::var("BLACK_BAG_VAULT_PATH") {
        let pb = PathBuf::from(path);
        if let Some(parent) = pb.parent() {
            fs::create_dir_all(parent)
                .with_context(|| format!("failed to create {}", parent.display()))?;
        }
        return Ok(pb);
    }
    let base = BaseDirs::new().ok_or_else(|| anyhow!("unable to resolve base directory"))?;
    let dir = base.config_dir().join("black_bag");
    fs::create_dir_all(&dir).with_context(|| format!("failed to create {}", dir.display()))?;
    Ok(dir.join("vault.cbor"))
}

fn prompt_passphrase(prompt: &str) -> Result<Zeroizing<String>> {
    let value = prompt_password(prompt)?;
    if value.trim().is_empty() {
        bail!("passphrase cannot be empty");
    }
    Ok(Zeroizing::new(value))
}

fn init_vault(cmd: InitCommand) -> Result<()> {
    let path = vault_path()?;
    let pass1 = prompt_passphrase("Master passphrase: ")?;
    let pass2 = prompt_passphrase("Confirm passphrase: ")?;
    if pass1.as_str() != pass2.as_str() {
        bail!("passphrases do not match");
    }
    Vault::init(&path, &pass1, cmd.mem_kib)?;
    println!("Initialized vault at {}", path.display());
    Ok(())
}

fn load_vault_with_prompt() -> Result<(Vault, Zeroizing<String>)> {
    let path = vault_path()?;
    let pass = prompt_passphrase("Master passphrase: ")?;
    let vault = Vault::load(&path, &pass)?;
    Ok((vault, pass))
}

fn add_record(cmd: AddCommand) -> Result<()> {
    let (mut vault, pass) = load_vault_with_prompt()?;

    let record = match cmd.record {
        AddRecord::Login(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let password = prompt_hidden("Password: ")?;
            Record::new(
                RecordData::Login {
                    username: args.username,
                    url: args.url,
                    password,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Contact(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            Record::new(
                RecordData::Contact {
                    full_name: args.full_name,
                    emails: args.emails,
                    phones: args.phones,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Id(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let secret = prompt_optional_hidden("Sensitive document secret (optional): ")?;
            Record::new(
                RecordData::Id {
                    id_type: args.id_type,
                    name_on_doc: args.name_on_doc,
                    number: args.number,
                    issuing_country: args.issuing_country,
                    expiry: args.expiry,
                    secret,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Note(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let body = prompt_multiline("Secure note body (Ctrl-D to finish): ")?;
            Record::new(RecordData::Note { body }, title, tags, notes)
        }
        AddRecord::Bank(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let account_number = prompt_hidden("Account number / secret: ")?;
            Record::new(
                RecordData::Bank {
                    institution: args.institution,
                    account_name: args.account_name,
                    routing_number: args.routing_number,
                    account_number,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Wifi(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let passphrase = prompt_hidden("Wi-Fi passphrase: ")?;
            Record::new(
                RecordData::Wifi {
                    ssid: args.ssid,
                    security: args.security,
                    location: args.location,
                    passphrase,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Api(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let secret = prompt_hidden("Secret key: ")?;
            Record::new(
                RecordData::Api {
                    service: args.service,
                    environment: args.environment,
                    access_key: args.access_key,
                    secret_key: secret,
                    scopes: args.scopes,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Wallet(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let secret = prompt_hidden("Wallet secret material: ")?;
            Record::new(
                RecordData::Wallet {
                    asset: args.asset,
                    address: args.address,
                    network: args.network,
                    secret_key: secret,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Totp(args) => {
            let AddTotp {
                common,
                issuer,
                account,
                secret,
                digits,
                step,
                skew,
                algorithm,
            } = args;
            if !(6..=8).contains(&digits) {
                bail!("digits must be between 6 and 8");
            }
            if step == 0 {
                bail!("step must be greater than zero");
            }
            let CommonRecordArgs { title, tags, notes } = common;
            let secret_bytes = match secret {
                Some(s) => parse_totp_secret(&s)?,
                None => {
                    let input = prompt_hidden("Base32 secret: ")?;
                    let value = Zeroizing::new(input.expose_utf8()?);
                    parse_totp_secret(value.as_str())?
                }
            };
            let totp_secret = Sensitive { data: secret_bytes };
            build_totp_instance(
                &totp_secret,
                digits,
                step,
                skew,
                algorithm,
                &issuer,
                &account,
            )?;
            Record::new(
                RecordData::Totp {
                    issuer,
                    account,
                    secret: totp_secret,
                    digits,
                    step,
                    skew,
                    algorithm,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Ssh(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let private_key = prompt_multiline_hidden("Paste private key (Ctrl-D to finish): ")?;
            Record::new(
                RecordData::Ssh {
                    label: args.label,
                    private_key,
                    comment: args.comment,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Pgp(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let armored =
                prompt_multiline_hidden("Paste armored private key (Ctrl-D to finish): ")?;
            Record::new(
                RecordData::Pgp {
                    label: args.label,
                    fingerprint: args.fingerprint,
                    armored_private_key: armored,
                },
                title,
                tags,
                notes,
            )
        }
        AddRecord::Recovery(args) => {
            let CommonRecordArgs { title, tags, notes } = args.common;
            let payload = prompt_multiline_hidden("Paste recovery payload (Ctrl-D to finish): ")?;
            Record::new(
                RecordData::Recovery {
                    description: args.description,
                    payload,
                },
                title,
                tags,
                notes,
            )
        }
    };

    vault.add_record(record);
    vault.save(&pass)?;
    println!("Record added");
    Ok(())
}
fn list_records(cmd: ListCommand) -> Result<()> {
    let (vault, _) = load_vault_with_prompt()?;
    let list = vault.list(cmd.kind, cmd.tag.as_deref(), cmd.query.as_deref());
    if list.is_empty() {
        println!("No matching records");
        return Ok(());
    }
    for record in list {
        println!(
            "{} | {} | {} | tags=[{}] | {}",
            record.id,
            record.kind(),
            record.title.as_deref().unwrap_or("(untitled)"),
            if record.tags.is_empty() {
                String::new()
            } else {
                record.tags.join(",")
            },
            record.data.summary_text()
        );
    }
    Ok(())
}

fn get_record(cmd: GetCommand) -> Result<()> {
    let (mut vault, _) = load_vault_with_prompt()?;
    if let Some(record) = vault.get(cmd.id) {
        println!("id: {}", record.id);
        println!("kind: {}", record.kind());
        if let Some(title) = &record.title {
            println!("title: {}", title);
        }
        if !record.tags.is_empty() {
            println!("tags: {}", record.tags.join(","));
        }
        if let Some(notes) = &record.metadata_notes {
            println!("notes: {}", notes);
        }
        if cmd.reveal {
            if !io::stdout().is_terminal() {
                bail!("--reveal requires an interactive TTY");
            }
            render_sensitive(record)?;
        } else {
            println!("(Sensitive fields hidden; re-run with --reveal on a TTY)");
        }
    } else {
        bail!("record {} not found", cmd.id);
    }
    Ok(())
}

fn render_sensitive(record: &Record) -> Result<()> {
    match &record.data {
        RecordData::Login { password, .. } => {
            println!("password: {}", password.expose_utf8()?);
        }
        RecordData::Contact { .. } => {}
        RecordData::Id { secret, .. } => {
            if let Some(secret) = secret {
                println!("secret: {}", secret.expose_utf8()?);
            }
        }
        RecordData::Note { body } => {
            println!("note:\n{}", body.expose_utf8()?);
        }
        RecordData::Bank { account_number, .. } => {
            println!("account_number: {}", account_number.expose_utf8()?);
        }
        RecordData::Wifi { passphrase, .. } => {
            println!("passphrase: {}", passphrase.expose_utf8()?);
        }
        RecordData::Api { secret_key, .. } => {
            println!("secret_key: {}", secret_key.expose_utf8()?);
        }
        RecordData::Wallet { secret_key, .. } => {
            println!("secret_key: {}", secret_key.expose_utf8()?);
        }
        RecordData::Totp {
            issuer,
            account,
            secret,
            digits,
            step,
            skew,
            algorithm,
        } => {
            let base32 = base32::encode(Rfc4648 { padding: false }, secret.as_slice());
            println!("secret_base32: {}", base32);
            if let Some(issuer) = issuer {
                println!("issuer: {}", issuer);
            }
            if let Some(account) = account {
                println!("account: {}", account);
            }
            println!("digits: {}", digits);
            println!("step: {}s", step);
            println!("skew: {}", skew);
            println!("algorithm: {}", algorithm);
            if let Ok(totp) =
                build_totp_instance(secret, *digits, *step, *skew, *algorithm, issuer, account)
            {
                if let Ok(code) = totp.generate_current() {
                    println!("current_code: {}", code);
                    if let Ok(ttl) = totp.ttl() {
                        println!("ttl: {}s", ttl);
                    }
                }
            }
        }
        RecordData::Ssh { private_key, .. } => {
            println!("private_key:\n{}", private_key.expose_utf8()?);
        }
        RecordData::Pgp {
            armored_private_key,
            ..
        } => {
            println!("pgp_private_key:\n{}", armored_private_key.expose_utf8()?);
        }
        RecordData::Recovery { payload, .. } => {
            println!("recovery_payload:\n{}", payload.expose_utf8()?);
        }
    }
    Ok(())
}

fn parse_totp_secret(input: &str) -> Result<Vec<u8>> {
    let cleaned: String = input
        .chars()
        .filter(|c| !c.is_whitespace() && *c != '-')
        .collect();
    if cleaned.is_empty() {
        bail!("secret cannot be empty");
    }
    let encoded = cleaned.to_uppercase();
    TotpSecret::Encoded(encoded)
        .to_bytes()
        .map_err(|_| anyhow!("invalid base32-encoded secret"))
}

fn build_totp_instance(
    secret: &Sensitive,
    digits: u8,
    step: u64,
    skew: u8,
    algorithm: TotpAlgorithm,
    issuer: &Option<String>,
    account: &Option<String>,
) -> Result<TOTP> {
    let account_name = account.clone().unwrap_or_default();
    TOTP::new(
        algorithm.to_lib(),
        usize::from(digits),
        skew,
        step,
        secret.as_slice().to_vec(),
        issuer.clone(),
        account_name,
    )
    .map_err(|err| anyhow!("failed to construct TOTP: {err}"))
}

fn split_secret(secret: &[u8], threshold: u8, share_count: u8) -> Result<Vec<Vec<u8>>> {
    if threshold == 0 {
        bail!("threshold must be at least 1");
    }
    if share_count < threshold {
        bail!("share count must be greater than or equal to threshold");
    }
    if share_count == 0 {
        bail!("share count must be at least 1");
    }
    let mut polys = Vec::with_capacity(secret.len());
    for &byte in secret {
        let mut poly = vec![byte];
        if threshold > 1 {
            let mut coeffs = vec![0u8; threshold.saturating_sub(1) as usize];
            OsRng.fill_bytes(&mut coeffs);
            poly.extend_from_slice(&coeffs);
        }
        polys.push(poly);
    }

    let mut shares = Vec::with_capacity(share_count as usize);
    for x in 1..=share_count {
        let mut share = Vec::with_capacity(secret.len() + 1);
        share.push(x);
        for poly in &polys {
            share.push(eval_poly(poly, x));
        }
        shares.push(share);
    }
    Ok(shares)
}

fn combine_secret(threshold: u8, shares: &[Vec<u8>]) -> Result<Vec<u8>> {
    if threshold == 0 {
        bail!("threshold must be at least 1");
    }
    if shares.len() < threshold as usize {
        bail!("not enough shares provided");
    }
    let first = shares
        .first()
        .ok_or_else(|| anyhow!("no shares provided"))?;
    if first.len() < 2 {
        bail!("share payload too short");
    }
    let share_len = first.len();
    for share in shares {
        if share.len() != share_len {
            bail!("shares have mismatched lengths");
        }
    }

    let mut seen = HashSet::new();
    let mut usable = Vec::new();
    for share in shares {
        let id = share[0];
        if id == 0 {
            bail!("share identifier must be non-zero");
        }
        if seen.insert(id) {
            usable.push((id, share[1..].to_vec()));
            if usable.len() == threshold as usize {
                break;
            }
        }
    }
    if usable.len() < threshold as usize {
        bail!("not enough unique shares to reconstruct secret");
    }

    let secret_len = share_len - 1;
    let mut secret = vec![0u8; secret_len];
    for idx in 0..secret_len {
        let points: Vec<(u8, u8)> = usable.iter().map(|(id, bytes)| (*id, bytes[idx])).collect();
        secret[idx] = interpolate_at_zero(&points)?;
    }
    Ok(secret)
}

fn eval_poly(poly: &[u8], x: u8) -> u8 {
    let mut acc = 0u8;
    for &coeff in poly.iter().rev() {
        acc = gf_mul(acc, x) ^ coeff;
    }
    acc
}

fn interpolate_at_zero(points: &[(u8, u8)]) -> Result<u8> {
    let mut result = 0u8;
    for (i, &(xi, yi)) in points.iter().enumerate() {
        let mut numerator = 1u8;
        let mut denominator = 1u8;
        for (j, &(xj, _)) in points.iter().enumerate() {
            if i == j {
                continue;
            }
            numerator = gf_mul(numerator, xj);
            let denom = xi ^ xj;
            if denom == 0 {
                bail!("duplicate share identifiers encountered");
            }
            denominator = gf_mul(denominator, denom);
        }
        let inv = gf_inv(denominator)?;
        let li = gf_mul(numerator, inv);
        result ^= gf_mul(yi, li);
    }
    Ok(result)
}

fn gf_mul(mut a: u8, mut b: u8) -> u8 {
    let mut p = 0u8;
    while b != 0 {
        if b & 1 != 0 {
            p ^= a;
        }
        let hi_bit = a & 0x80;
        a <<= 1;
        if hi_bit != 0 {
            a ^= 0x1b;
        }
        b >>= 1;
    }
    p
}

fn gf_pow(mut base: u8, mut exp: u16) -> u8 {
    let mut result = 1u8;
    while exp != 0 {
        if exp & 1 != 0 {
            result = gf_mul(result, base);
        }
        base = gf_mul(base, base);
        exp >>= 1;
    }
    result
}

fn gf_inv(x: u8) -> Result<u8> {
    if x == 0 {
        bail!("encountered zero denominator during interpolation");
    }
    Ok(gf_pow(x, 254))
}

fn rotate_vault(cmd: RotateCommand) -> Result<()> {
    let (mut vault, pass) = load_vault_with_prompt()?;
    vault.rotate(&pass, cmd.mem_kib)?;
    vault.save(&pass)?;
    println!("Rotation complete");
    Ok(())
}

fn doctor(cmd: DoctorCommand) -> Result<()> {
    let (vault, _) = load_vault_with_prompt()?;
    let stats = vault.stats();
    if cmd.json {
        let payload = json!({
            "ready": true,
            "recordCount": stats.record_count,
            "argonMemKib": stats.argon_mem_kib,
            "argonTimeCost": stats.argon_time_cost,
            "argonLanes": stats.argon_lanes,
            "createdAt": stats.created_at.to_rfc3339(),
            "updatedAt": stats.updated_at.to_rfc3339(),
        });
        println!("{}", payload);
    } else {
        println!("status: ready");
        println!("records: {}", stats.record_count);
        println!("created: {}", stats.created_at);
        println!("updated: {}", stats.updated_at);
        println!(
            "argon2: mem={} KiB, time={}, lanes={}",
            stats.argon_mem_kib, stats.argon_time_cost, stats.argon_lanes
        );
    }
    Ok(())
}

fn recovery(cmd: RecoveryCommand) -> Result<()> {
    match cmd {
        RecoveryCommand::Split(args) => {
            if args.threshold == 0 || args.threshold > args.shares {
                bail!("threshold must be between 1 and number of shares");
            }
            let secret = prompt_hidden("Secret to split: ")?;
            let shares = split_secret(secret.as_slice(), args.threshold, args.shares)?;
            for share in shares {
                let (id, data) = share
                    .split_first()
                    .ok_or_else(|| anyhow!("invalid share structure"))?;
                let encoded = BASE64.encode(data);
                println!("{}-{}", id, encoded);
            }
            Ok(())
        }
        RecoveryCommand::Combine(args) => {
            if args.threshold == 0 {
                bail!("threshold must be at least 1");
            }
            let mut shares = Vec::new();
            for part in args
                .shares
                .split(',')
                .map(str::trim)
                .filter(|s| !s.is_empty())
            {
                let (id, data) = part
                    .split_once('-')
                    .ok_or_else(|| anyhow!("invalid share format: {part}"))?;
                let identifier: u8 = id.parse().context("invalid share identifier")?;
                if identifier == 0 {
                    bail!("share identifier must be between 1 and 255");
                }
                let mut decoded = BASE64.decode(data).context("invalid base64 in share")?;
                if decoded.is_empty() {
                    bail!("share payload cannot be empty");
                }
                let mut share = Vec::with_capacity(decoded.len() + 1);
                share.push(identifier);
                share.append(&mut decoded);
                shares.push(share);
            }
            if shares.len() < args.threshold as usize {
                bail!(
                    "insufficient shares provided (need at least {})",
                    args.threshold
                );
            }
            let secret = combine_secret(args.threshold, &shares)?;
            println!("{}", String::from_utf8_lossy(&secret));
            Ok(())
        }
    }
}

fn totp(cmd: TotpCommand) -> Result<()> {
    match cmd {
        TotpCommand::Code(args) => totp_code(args),
    }
}

fn totp_code(args: TotpCodeCommand) -> Result<()> {
    let (vault, _) = load_vault_with_prompt()?;
    let record = vault
        .get_ref(args.id)
        .ok_or_else(|| anyhow!("record {} not found", args.id))?;
    let (issuer, account, secret, digits, step, skew, algorithm) = match &record.data {
        RecordData::Totp {
            issuer,
            account,
            secret,
            digits,
            step,
            skew,
            algorithm,
        } => (issuer, account, secret, *digits, *step, *skew, *algorithm),
        _ => bail!("record {} is not a TOTP secret", args.id),
    };

    let totp = build_totp_instance(secret, digits, step, skew, algorithm, issuer, account)?;
    let code = if let Some(ts) = args.time {
        if ts < 0 {
            bail!("time must be non-negative");
        }
        totp.generate(ts as u64)
    } else {
        totp.generate_current()?
    };
    println!("code: {}", code);
    if args.time.is_none() {
        let ttl = totp.ttl()?;
        println!("ttl: {}s", ttl);
    }
    Ok(())
}

fn self_test() -> Result<()> {
    let mut sample = [0u8; 32];
    OsRng.fill_bytes(&mut sample);
    let secret = Sensitive::new_from_utf8(&sample);
    let record = Record::new(
        RecordData::Note { body: secret },
        Some("Self-test".into()),
        vec!["selftest".into()],
        None,
    );
    let payload = VaultPayload {
        records: vec![record],
        record_counter: 1,
    };

    let mut dek = [0u8; 32];
    OsRng.fill_bytes(&mut dek);
    let blob = encrypt_payload(&dek, &payload)?;
    let recovered = decrypt_payload(&dek, &blob)?;
    anyhow::ensure!(recovered.records.len() == 1, "self-test failed");
    println!("Self-test passed");
    Ok(())
}

fn prompt_hidden(prompt: &str) -> Result<Sensitive> {
    let value = Zeroizing::new(prompt_password(prompt)?);
    Ok(Sensitive::from_string(value.as_str()))
}

fn prompt_optional_hidden(prompt: &str) -> Result<Option<Sensitive>> {
    let value = Zeroizing::new(prompt_password(prompt)?);
    if value.trim().is_empty() {
        Ok(None)
    } else {
        Ok(Some(Sensitive::from_string(value.as_str())))
    }
}

fn prompt_multiline(prompt: &str) -> Result<Sensitive> {
    eprintln!("{}", prompt);
    read_multiline(false)
}

fn prompt_multiline_hidden(prompt: &str) -> Result<Sensitive> {
    eprintln!("{}", prompt);
    read_multiline(true)
}

fn read_multiline(_hidden: bool) -> Result<Sensitive> {
    let mut buffer = Zeroizing::new(Vec::new());
    io::stdin().read_to_end(&mut buffer)?;
    while buffer.last().copied() == Some(b'\n') {
        buffer.pop();
    }
    Ok(Sensitive::new_from_utf8(&buffer))
}

#[cfg(feature = "fuzzing")]
pub fn fuzz_try_payload(bytes: &[u8]) {
    let _ = ciborium::de::from_reader::<VaultPayload, _>(bytes);
}

fn write_vault(path: &Path, file: &VaultFile) -> Result<()> {
    let parent = path.parent().ok_or_else(|| anyhow!("invalid vault path"))?;
    fs::create_dir_all(parent)?;
    let mut tmp = NamedTempFile::new_in(parent)?;
    into_writer(file, &mut tmp).context("failed to serialize vault")?;
    tmp.as_file_mut().sync_all()?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        tmp.as_file_mut()
            .set_permissions(fs::Permissions::from_mode(0o600))?;
    }
    tmp.persist(path)?;
    Ok(())
}

impl fmt::Display for RecordKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let label = match self {
            RecordKind::Login => "login",
            RecordKind::Contact => "contact",
            RecordKind::Id => "id",
            RecordKind::Note => "note",
            RecordKind::Bank => "bank",
            RecordKind::Wifi => "wifi",
            RecordKind::Api => "api",
            RecordKind::Wallet => "wallet",
            RecordKind::Totp => "totp",
            RecordKind::Ssh => "ssh",
            RecordKind::Pgp => "pgp",
            RecordKind::Recovery => "recovery",
        };
        f.write_str(label)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;
    use proptest::proptest;
    use proptest::strategy::Strategy;
    use serial_test::serial;
    use std::env;
    use std::path::PathBuf;
    use tempfile::tempdir;

    fn prepare(passphrase: &str) -> Result<(tempfile::TempDir, PathBuf, Zeroizing<String>)> {
        let dir = tempdir()?;
        let vault_path = dir.path().join("vault.cbor");
        env::set_var("BLACK_BAG_VAULT_PATH", &vault_path);
        let pass = Zeroizing::new(passphrase.to_string());
        Ok((dir, vault_path, pass))
    }

    fn cleanup() {
        env::remove_var("BLACK_BAG_VAULT_PATH");
    }

    fn arb_ascii_string(max: usize) -> impl Strategy<Value = String> {
        proptest::collection::vec(proptest::char::range('a', 'z'), 0..=max)
            .prop_map(|chars| chars.into_iter().collect())
    }

    fn arb_note_record() -> impl Strategy<Value = Record> {
        (
            proptest::option::of(arb_ascii_string(12)),
            proptest::collection::vec(arb_ascii_string(8), 0..3),
            arb_ascii_string(48),
        )
            .prop_map(|(title, tags, body)| {
                Record::new(
                    RecordData::Note {
                        body: Sensitive::from_string(&body),
                    },
                    title,
                    tags,
                    None,
                )
            })
    }

    proptest! {
        #[test]
        fn encrypt_blob_roundtrip_prop(
            key_bytes in proptest::array::uniform32(any::<u8>()),
            data in proptest::collection::vec(any::<u8>(), 0..256),
            aad in proptest::collection::vec(any::<u8>(), 0..32),
        ) {
            let blob = encrypt_blob(&key_bytes, &data, &aad).unwrap();
            let decrypted = decrypt_blob(&key_bytes, &blob, &aad).unwrap();
            prop_assert_eq!(decrypted, data);
        }

        #[test]
        fn payload_roundtrip_prop(
            key_bytes in proptest::array::uniform32(any::<u8>()),
            records in proptest::collection::vec(arb_note_record(), 0..3),
        ) {
            let payload = VaultPayload {
                records: records.clone(),
                record_counter: records.len() as u64,
            };
            let blob = encrypt_payload(&key_bytes, &payload).unwrap();
            let decoded = decrypt_payload(&key_bytes, &blob).unwrap();
            prop_assert_eq!(decoded, payload);
        }
    }

    #[test]
    #[serial]
    fn vault_round_trip_note() -> Result<()> {
        let (_tmp, vault_path, pass) = prepare("correct horse battery staple")?;
        Vault::init(&vault_path, &pass, 32_768)?;

        let mut vault = Vault::load(&vault_path, &pass)?;
        let record = Record::new(
            RecordData::Note {
                body: Sensitive::from_string("mission ops"),
            },
            Some("Ops Note".into()),
            vec!["mission".into()],
            None,
        );
        let record_id = record.id;
        vault.add_record(record);
        vault.save(&pass)?;

        drop(vault);
        let vault = Vault::load(&vault_path, &pass)?;
        let notes = vault.list(Some(RecordKind::Note), None, None);
        assert_eq!(notes.len(), 1);
        assert_eq!(notes[0].id, record_id);
        assert_eq!(notes[0].title.as_deref(), Some("Ops Note"));
        assert!(notes[0].matches_tag("mission"));

        cleanup();
        Ok(())
    }

    #[test]
    #[serial]
    fn totp_round_trip() -> Result<()> {
        let (_tmp, vault_path, pass) = prepare("totp-pass")?;
        Vault::init(&vault_path, &pass, 32_768)?;

        let secret_bytes = parse_totp_secret("JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP")?;
        let mut vault = Vault::load(&vault_path, &pass)?;
        let record = Record::new(
            RecordData::Totp {
                issuer: Some("TestIssuer".into()),
                account: Some("test@example".into()),
                secret: Sensitive { data: secret_bytes },
                digits: 6,
                step: 30,
                skew: 1,
                algorithm: TotpAlgorithm::Sha1,
            },
            Some("TOTP".into()),
            vec![],
            None,
        );
        let record_id = record.id;
        vault.add_record(record);
        vault.save(&pass)?;

        drop(vault);
        let vault = Vault::load(&vault_path, &pass)?;
        let record = vault
            .get_ref(record_id)
            .ok_or_else(|| anyhow!("TOTP record missing"))?;
        let code = match &record.data {
            RecordData::Totp {
                issuer,
                account,
                secret,
                digits,
                step,
                skew,
                algorithm,
            } => build_totp_instance(secret, *digits, *step, *skew, *algorithm, issuer, account)?
                .generate(59),
            _ => bail!("expected totp record"),
        };
        let expected = TOTP::new(
            TotpAlgorithmLib::SHA1,
            6,
            1,
            30,
            parse_totp_secret("JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP")?,
            Some("TestIssuer".into()),
            "test@example".into(),
        )
        .unwrap()
        .generate(59);
        assert_eq!(code, expected);

        cleanup();
        Ok(())
    }

    #[test]
    #[serial]
    fn vault_rotation_changes_wrapped_keys() -> Result<()> {
        let (_tmp, vault_path, pass) = prepare("rotate-all-the-things")?;
        Vault::init(&vault_path, &pass, 32_768)?;

        let mut vault = Vault::load(&vault_path, &pass)?;
        let record = Record::new(
            RecordData::Api {
                service: Some("intel-api".into()),
                environment: Some("prod".into()),
                access_key: Some("AKIA-123".into()),
                secret_key: Sensitive::from_string("super-secret"),
                scopes: vec!["read".into(), "write".into()],
            },
            Some("API".into()),
            vec!["read".into()],
            None,
        );
        vault.add_record(record);
        let before = vault.file.header.sealed_dek.ciphertext.clone();
        vault.rotate(&pass, Some(65_536))?;
        vault.save(&pass)?;
        let after = vault.file.header.sealed_dek.ciphertext.clone();
        assert_ne!(before, after);

        drop(vault);
        let vault = Vault::load(&vault_path, &pass)?;
        let apis = vault.list(Some(RecordKind::Api), Some("read"), None);
        assert_eq!(apis.len(), 1);
        assert!(apis[0].data.summary_text().contains("intel-api"));

        cleanup();
        Ok(())
    }

    #[test]
    fn recovery_split_combine_roundtrip() -> Result<()> {
        let secret = b"ultra-secret";
        let shares = split_secret(secret, 3, 5)?;
        let recovered = combine_secret(3, &shares)?;
        assert_eq!(recovered, secret);
        Ok(())
    }
}
