//! implementation of the `vanity` subcommand

use clap::Parser;
use rand::{rngs::OsRng, RngCore};
use sc_cli::{utils::format_seed, with_crypto_scheme, CryptoSchemeFlag, Error, OutputTypeFlag};
use sp_core::crypto::{unwrap_or_default_ss58_version, Ss58AddressFormat, Ss58Codec};
use sp_runtime::traits::IdentifyAccount;

use crate::commands::utils::print_from_uri;
use crate::params::NetworkSchemeFlag;

/// The `vanity` command
#[derive(Debug, Clone, Parser)]
#[command(
    name = "vanity",
    about = "Generate a seed that provides a vanity address"
)]
pub struct VanityCmd {
    /// Desired pattern
    #[arg(long, value_parser = assert_non_empty_string)]
    pattern: String,

    #[allow(missing_docs)]
    #[clap(flatten)]
    network_scheme: NetworkSchemeFlag,

    #[allow(missing_docs)]
    #[clap(flatten)]
    output_scheme: OutputTypeFlag,

    #[allow(missing_docs)]
    #[clap(flatten)]
    crypto_scheme: CryptoSchemeFlag,
}

impl VanityCmd {
    /// Run the command
    pub fn run(&self) -> Result<(), Error> {
        let formated_seed = with_crypto_scheme!(
            self.crypto_scheme.scheme,
            generate_key(
                &self.pattern,
                unwrap_or_default_ss58_version(self.network_scheme.network)
            ),
        )?;

        with_crypto_scheme!(
            self.crypto_scheme.scheme,
            print_from_uri(
                &formated_seed,
                None,
                self.network_scheme.network,
                self.output_scheme.output_type,
            ),
        );
        Ok(())
    }
}

/// genertae a key based on given pattern
fn generate_key<Pair>(
    desired: &str,
    network_override: Ss58AddressFormat,
) -> Result<String, &'static str>
where
    Pair: sp_core::Pair,
    Pair::Public: IdentifyAccount,
    <Pair::Public as IdentifyAccount>::AccountId: Ss58Codec,
{
    println!("Generating key containing pattern '{}'", desired);

    let top = 45 + (desired.len() * 48);
    let mut best = 0;
    let mut seed = Pair::Seed::default();
    let mut done = 0;

    loop {
        if done % 100000 == 0 {
            OsRng.fill_bytes(seed.as_mut());
        } else {
            next_seed(seed.as_mut());
        }

        let p = Pair::from_seed(&seed);
        let ss58 = p
            .public()
            .into_account()
            .to_ss58check_with_version(network_override);
        println!("{:?}", ss58);
        let score = calculate_score(desired, &ss58);
        if score > best || desired.len() < 2 {
            best = score;
            if best >= top {
                println!("best: {} == top: {}", best, top);
                return Ok(format_seed::<Pair>(seed.clone()));
            }
        }
        done += 1;

        if done % good_waypoint(done) == 0 {
            println!("{} keys searched; best is {}/{} complete", done, best, top);
        }
    }
}

fn good_waypoint(done: u64) -> u64 {
    match done {
        0..=1_000_000 => 100_000,
        1_000_001..=10_000_000 => 1_000_000,
        10_000_001..=100_000_000 => 10_000_000,
        100_000_001.. => 100_000_000,
    }
}

fn next_seed(seed: &mut [u8]) {
    for s in seed {
        match s {
            255 => {
                *s = 0;
            }
            _ => {
                *s += 1;
                break;
            }
        }
    }
}

/// Calculate the score of a key based on the desired
/// input.
fn calculate_score(_desired: &str, key: &str) -> usize {
    for truncate in 0.._desired.len() {
        let snip_size = _desired.len() - truncate;
        let truncated = &_desired[0..snip_size];
        if let Some(pos) = key.find(truncated) {
            return (47 - pos) + (snip_size * 48);
        }
    }
    0
}

/// checks that `pattern` is non-empty
fn assert_non_empty_string(pattern: &str) -> Result<String, &'static str> {
    if pattern.is_empty() {
        Err("Pattern must not be empty")
    } else {
        Ok(pattern.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use sp_core::{
        crypto::{default_ss58_version, Ss58AddressFormatRegistry, Ss58Codec},
        sr25519, Pair,
    };

    #[test]
    fn vanity() {
        let vanity = VanityCmd::parse_from(&["vanity", "--pattern", "j"]);
        assert!(vanity.run().is_ok());
    }

    #[test]
    fn test_generation_with_single_char() {
        let seed = generate_key::<sr25519::Pair>("ab", default_ss58_version()).unwrap();
        assert!(
            sr25519::Pair::from_seed_slice(&array_bytes::hex2bytes_unchecked(&seed))
                .unwrap()
                .public()
                .to_ss58check()
                .contains("ab")
        );
    }

    #[test]
    fn generate_key_respects_network_override() {
        let seed =
            generate_key::<sr25519::Pair>("ab", Ss58AddressFormatRegistry::PolkadotAccount.into())
                .unwrap();
        assert!(
            sr25519::Pair::from_seed_slice(&array_bytes::hex2bytes_unchecked(&seed))
                .unwrap()
                .public()
                .to_ss58check_with_version(Ss58AddressFormatRegistry::PolkadotAccount.into())
                .contains("ab")
        );
    }

    #[test]
    fn test_score_1_char_100() {
        let score = calculate_score("j", "sYrjWiZkEP1PQe2kKEZiC1Bi9L8yFSsB5RPEkzEPd5NThsB5H");
        assert_eq!(score, 96);
    }

    #[test]
    fn test_score_100() {
        let score = calculate_score("ghst", "sYsghstuhYd9p6unUC3kPxjD2gv2zRCztYQaEDCMJpYrPTqTG");
        assert_eq!(score, 246);
    }

    #[test]
    fn test_score_50_2() {
        // 50% for the position + 50% for the size
        assert_eq!(
            calculate_score("ghst", "sYsghXXuhYd9p6unUC3kPxjD2gv2zRCztYQaEDCMJpYrPTqTG"),
            146
        );
    }

    #[test]
    fn test_score_0() {
        assert_eq!(
            calculate_score("ghst", "sYrjWiZkEP1PQe2kKEZiC1Bi9L8yFSsB5RPEkzEPd5NThsB5H"),
            0
        );
    }
}
