// Copyright The Hush Developers  2019-2022
// Released under the GPLv3

//! The Sapling params on a `LightClient` borrow the image that rust-embed bakes into the binary
//! (`Cow::Borrowed(&'static [u8])`) instead of copying ~49 MiB onto the heap for every client.
//!
//! Every client here is on the "main" chain (the "test" chain, which `LightClient::unconnected`
//! hard-codes, panics in `base58_pubkey_address` when the wallet is created) and lives in a temp
//! data dir `<tmp>/data`, so the wallet file and `get_zcash_params_path()` (`<data>/../.zcash-params`)
//! both stay inside the temp dir: nothing is written to ~/.silentdragonxlite or ~/.zcash-params.

use std::borrow::Cow;
use std::path::PathBuf;

use bellman::groth16::Parameters;
use pairing::bls12_381::Bls12;
use sha2::{Digest, Sha256};
use tempdir::TempDir;

use silentdragonxlitelib::lightclient::{LightClient, LightClientConfig};
use silentdragonxlitelib::SaplingParams;

const SEED: &str = "youth strong sweet gorilla hammer unhappy congress stamp left stereo riot salute road tag clean toilet artefact fork certain leopard entire civil degree wonder";

const OUTPUT_LEN: usize = 3_592_860;
const SPEND_LEN: usize = 47_958_396;
const OUTPUT_SHA256: &str = "2f0ebbcbb9bb0bcffe95a397e7eba89c29eb4dde6191c339db88570e3f3fb0e4";
const SPEND_SHA256: &str = "8e48ffd23abb3a5fd9c5589204f32d9c31285a04b78096ba40a79b75677efc13";

fn main_config(tmp: &TempDir) -> LightClientConfig {
    let data: PathBuf = tmp.path().join("data");
    LightClientConfig::create_unconnected("main".to_string(), data.to_str().map(|s| s.to_string()))
}

fn embedded(name: &str) -> Cow<'static, [u8]> {
    SaplingParams::get(name).expect("embedded Sapling params")
}

fn sha256_hex(bytes: &[u8]) -> String {
    hex::encode(Sha256::digest(bytes))
}

fn is_borrowed(c: &Cow<'static, [u8]>) -> bool {
    matches!(c, Cow::Borrowed(_))
}

/// Borrowed (no heap copy), the right length and content, and the files are not swapped.
fn assert_params_pinned(lc: &LightClient, ctx: &str) {
    assert!(is_borrowed(&lc.sapling_output), "{}: sapling_output is not Cow::Borrowed", ctx);
    assert!(is_borrowed(&lc.sapling_spend), "{}: sapling_spend is not Cow::Borrowed", ctx);
    assert_eq!(lc.sapling_output.len(), OUTPUT_LEN, "{}: sapling_output length", ctx);
    assert_eq!(lc.sapling_spend.len(), SPEND_LEN, "{}: sapling_spend length", ctx);
    assert_eq!(sha256_hex(&lc.sapling_output[..]), OUTPUT_SHA256, "{}: sapling_output hash", ctx);
    assert_eq!(sha256_hex(&lc.sapling_spend[..]), SPEND_SHA256, "{}: sapling_spend hash", ctx);
}

#[test]
fn params_borrowed_and_pinned() {
    // Restore from a seed phrase.
    let tmp = TempDir::new("sdxl-params").unwrap();
    let cfg = main_config(&tmp);
    let lc = LightClient::new_from_phrase(SEED.to_string(), &cfg, 0, 0, false).unwrap();
    assert_params_pinned(&lc, "new_from_phrase");
    // The borrow points into the embedded image itself.
    assert_eq!(lc.sapling_spend.as_ptr(), embedded("sapling-spend.params").as_ptr());
    assert_eq!(lc.sapling_output.as_ptr(), embedded("sapling-output.params").as_ptr());

    // The GUI open path: create a new wallet, save it, then open it from disk.
    let tmp2 = TempDir::new("sdxl-params").unwrap();
    let cfg2 = main_config(&tmp2);
    let lc2 = LightClient::new(&cfg2, 0).unwrap();
    assert_params_pinned(&lc2, "new");
    lc2.do_save().unwrap();
    drop(lc2);
    let lc3 = LightClient::read_from_disk(&cfg2).unwrap();
    assert_params_pinned(&lc3, "read_from_disk");
}

#[test]
fn set_params_guard_and_hash() {
    let tmp = TempDir::new("sdxl-params").unwrap();
    let cfg = main_config(&tmp);
    let mut lc = LightClient::new_from_phrase(SEED.to_string(), &cfg, 0, 0, false).unwrap();

    let out = embedded("sapling-output.params").to_vec();
    let spend = embedded("sapling-spend.params").to_vec();

    // Fields already filled from the embedded image: a valid set is Ok and does not overwrite them.
    let p_out = lc.sapling_output.as_ptr();
    let p_spend = lc.sapling_spend.as_ptr();
    assert!(lc.set_sapling_params(&out, &spend).is_ok());
    assert!(is_borrowed(&lc.sapling_output) && is_borrowed(&lc.sapling_spend));
    assert_eq!(lc.sapling_output.as_ptr(), p_out);
    assert_eq!(lc.sapling_spend.as_ptr(), p_spend);

    // Empty fields (as with embed_params off).
    lc.sapling_output = Cow::Borrowed(&[]);
    lc.sapling_spend = Cow::Borrowed(&[]);

    // Swapped files fail the hash check and leave both fields untouched.
    assert!(lc.set_sapling_params(&spend, &out).is_err());
    assert!(lc.sapling_output.is_empty());
    assert!(lc.sapling_spend.is_empty());

    // The right files are taken as owned copies with identical bytes.
    assert!(lc.set_sapling_params(&out, &spend).is_ok());
    assert!(matches!(lc.sapling_output, Cow::Owned(_)), "sapling_output is not Cow::Owned");
    assert!(matches!(lc.sapling_spend, Cow::Owned(_)), "sapling_spend is not Cow::Owned");
    assert!(lc.sapling_output[..] == out[..]);
    assert!(lc.sapling_spend[..] == spend[..]);

    // A second call does not overwrite them.
    let p = lc.sapling_spend.as_ptr();
    let po = lc.sapling_output.as_ptr();
    assert!(lc.set_sapling_params(&out, &spend).is_ok());
    assert_eq!(lc.sapling_spend.as_ptr(), p);
    assert_eq!(lc.sapling_output.as_ptr(), po);

    // The params are also written to <data>/../.zcash-params, i.e. inside the temp dir.
    #[cfg(target_os = "linux")]
    {
        let dir = tmp.path().join(".zcash-params");
        let spend_file = dir.join("sapling-spend.params");
        let output_file = dir.join("sapling-output.params");
        assert!(spend_file.exists(), "{} missing", spend_file.display());
        assert!(output_file.exists(), "{} missing", output_file.display());
        assert_eq!(std::fs::metadata(&spend_file).unwrap().len(), SPEND_LEN as u64);
        assert_eq!(std::fs::metadata(&output_file).unwrap().len(), OUTPUT_LEN as u64);
    }
}

#[test]
fn borrowed_params_parse_identically() {
    let tmp = TempDir::new("sdxl-params").unwrap();
    let cfg = main_config(&tmp);
    let lc = LightClient::new_from_phrase(SEED.to_string(), &cfg, 0, 0, false).unwrap();

    // The new borrowed path against the old heap-copy path: the prover gets identical keys.
    let spend_new = Parameters::<Bls12>::read(&lc.sapling_spend[..], false).unwrap();
    let spend_old = Parameters::<Bls12>::read(&embedded("sapling-spend.params").to_vec()[..], false).unwrap();
    assert!(spend_new == spend_old, "spend params parse differently");
    drop(spend_new);
    drop(spend_old);

    let output_new = Parameters::<Bls12>::read(&lc.sapling_output[..], false).unwrap();
    let output_old = Parameters::<Bls12>::read(&embedded("sapling-output.params").to_vec()[..], false).unwrap();
    assert!(output_new == output_old, "output params parse differently");
}

#[cfg(target_os = "linux")]
fn rss_anon_kib() -> u64 {
    let status = std::fs::read_to_string("/proc/self/status").unwrap();
    for line in status.lines() {
        if line.starts_with("RssAnon:") {
            return line["RssAnon:".len()..].trim().trim_end_matches("kB").trim().parse().unwrap();
        }
    }
    panic!("no RssAnon in /proc/self/status");
}

/// Run alone: `--test sapling_params -- --ignored --test-threads=1`. The old code copied both
/// params onto the heap for every client (~49 MiB of RssAnon).
#[cfg(target_os = "linux")]
#[test]
#[ignore]
fn open_rss_anon_delta() {
    // Warm up one-time allocations (lazy statics, allocator arenas) with a first client.
    let tmp0 = TempDir::new("sdxl-params").unwrap();
    let _warm = LightClient::new_from_phrase(SEED.to_string(), &main_config(&tmp0), 0, 0, false).unwrap();

    let tmp = TempDir::new("sdxl-params").unwrap();
    let cfg = main_config(&tmp);
    let before = rss_anon_kib();
    let lc = LightClient::new_from_phrase(SEED.to_string(), &cfg, 0, 0, false).unwrap();
    let after = rss_anon_kib();
    let delta = after.saturating_sub(before);
    println!("RssAnon before {} kB, after {} kB, delta {} kB", before, after, delta);
    assert_eq!(lc.sapling_spend.len(), SPEND_LEN);
    assert!(delta < 16 * 1024, "opening a client grew RssAnon by {} kB", delta);
}