pub struct Pubkey(pub(crate) [u8; 32]);

pub fn find_program_address(seeds: &[&[u8]], program_id: &Pubkey) -> (Pubkey, u8) {
    Self::try_find_program_address(seeds, program_id)
        .unwrap_or_else(|| panic!("Unable to find a viable program address bump seed"))
}

pub fn try_find_program_address(seeds: &[&[u8]], program_id: &Pubkey) -> Option<(Pubkey, u8)> {
    // Perform the calculation inline, calling this from within a program is
    // not supported
    #[cfg(not(target_os = "solana"))]
    {
        let mut bump_seed = [std::u8::MAX]; // 从 std::u8::MAX 开始尝试
        for _ in 0..std::u8::MAX {
            {
                let mut seeds_with_bump = seeds.to_vec();
                seeds_with_bump.push(&bump_seed); // 将 bump 作为种子的最后一部分
                match Self::create_program_address(&seeds_with_bump, program_id) {
                    Ok(address) => return Some((address, bump_seed[0])),
                    Err(PubkeyError::InvalidSeeds) => (),
                    _ => break,
                }
            }
            bump_seed[0] -= 1; // 尝试失败了就减 1 再尝试
        }
        None
    }
    // Call via a system call to perform the calculation
    #[cfg(target_os = "solana")]
    {
        // ...
    }
}

pub fn create_program_address(
    seeds: &[&[u8]],
    program_id: &Pubkey,
) -> Result<Pubkey, PubkeyError> {
    if seeds.len() > MAX_SEEDS { // 种子最多有 16 个
        return Err(PubkeyError::MaxSeedLengthExceeded);
    }
    for seed in seeds.iter() {
        if seed.len() > MAX_SEED_LEN { // 每个种子最长 32
            return Err(PubkeyError::MaxSeedLengthExceeded);
        }
    }

    // Perform the calculation inline, calling this from within a program is
    // not supported
    #[cfg(not(target_os = "solana"))]
    {
        let mut hasher = crate::hash::Hasher::default(); // 就是一个 Sha256
        for seed in seeds.iter() {
            hasher.hash(seed);
        }
        // PDA_MARKER 是一个 21 字节长的字符串
        // const PDA_MARKER: &[u8; 21] = b"ProgramDerivedAddress";
        hasher.hashv(&[program_id.as_ref(), PDA_MARKER]);
        let hash = hasher.result();

        if bytes_are_curve_point(hash) { // PDA 账户需要确保不在爱德华曲线上
            return Err(PubkeyError::InvalidSeeds);
        }

        Ok(Pubkey::from(hash.to_bytes()))
    }
    // Call via a system call to perform the calculation
    #[cfg(target_os = "solana")]
    {
        let mut bytes = [0; 32];
        let result = unsafe {
            crate::syscalls::sol_create_program_address(
                seeds as *const _ as *const u8,
                seeds.len() as u64,
                program_id as *const _ as *const u8,
                &mut bytes as *mut _ as *mut u8,
            )
        };
        match result {
            crate::entrypoint::SUCCESS => Ok(Pubkey::from(bytes)),
            _ => Err(result.into()),
        }
    }
}

pub fn bytes_are_curve_point<T: AsRef<[u8]>>(_bytes: T) -> bool {
    #[cfg(not(target_os = "solana"))]
    {
        curve25519_dalek::edwards::CompressedEdwardsY::from_slice(_bytes.as_ref())
            .decompress()
            .is_some()
    }
    #[cfg(target_os = "solana")]
    unimplemented!();
}

/// Account information
#[derive(Clone)]
#[repr(C)]
pub struct AccountInfo<'a> {
    /// Public key of the account
    pub key: &'a Pubkey,
    /// The lamports in the account.  Modifiable by programs.
    pub lamports: Rc<RefCell<&'a mut u64>>,
    /// The data held in this account.  Modifiable by programs.
    pub data: Rc<RefCell<&'a mut [u8]>>,
    /// Program that owns this account
    pub owner: &'a Pubkey,
    /// The epoch at which this account will next owe rent
    pub rent_epoch: Epoch,
    /// Was the transaction signed by this account's public key?
    pub is_signer: bool,
    /// Is the account writable?
    pub is_writable: bool,
    /// This account's data contains a loaded program (and is now read-only)
    pub executable: bool,
}

#[derive(Clone)]
pub struct Account<'info, T: AccountSerialize + AccountDeserialize + Clone> {
    account: T,
    info: &'info AccountInfo<'info>,
}

#[derive(Clone)]
pub struct AccountLoader<'info, T: ZeroCopy + Owner> {
    acc_info: &'info AccountInfo<'info>,
    phantom: PhantomData<&'info T>,
}

#[derive(Clone)]
pub struct InterfaceAccount<'info, T: AccountSerialize + AccountDeserialize + Clone> {
    account: Account<'info, T>,
    // The owner here is used to make sure that changes aren't incorrectly propagated
    // to an account with a modified owner
    owner: Pubkey,
}

#[derive(Clone)]
pub struct Interface<'info, T>(Program<'info, T>);

#[derive(Clone)]
pub struct Program<'info, T> {
    info: &'info AccountInfo<'info>,
    _phantom: PhantomData<T>,
}

#[derive(Debug, Clone)]
pub struct Signer<'info> {
    info: &'info AccountInfo<'info>,
}

#[derive(Debug, Clone)]
pub struct SystemAccount<'info> {
    info: &'info AccountInfo<'info>,
}

pub struct Sysvar<'info, T: solana_program::sysvar::Sysvar> {
    info: &'info AccountInfo<'info>,
    account: T,
}

#[derive(Debug, Clone)]
pub struct UncheckedAccount<'info>(&'info AccountInfo<'info>);

use num::BigInt;
use rand::seq::index::sample;
use rand::{thread_rng, Rng};

const SMALL_PRIMES: [u64; 50] = [
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97,
    101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193,
    197, 199, 211, 223, 227, 229,
];

#[derive(Debug)]
struct KeyPair {
    private_key: u64,
    public_key: u64,
    modulus: u64,
}

impl KeyPair {
    pub fn generate() -> Self {
        let mut rng = thread_rng();
        let indices = sample(&mut rng, SMALL_PRIMES.len(), 2);
        // use crate rand to select two random prime numbers
        let p = SMALL_PRIMES[indices.index(0)];
        let q = SMALL_PRIMES[indices.index(1)];
        let modulus = p * q;
        let phi = (p - 1) * (q - 1);
        let public_key = loop {
            let public_key = SMALL_PRIMES[thread_rng().gen_range(0..SMALL_PRIMES.len())];
            if public_key < phi && gcd(public_key, phi) == 1 {
                break public_key;
            }
        };

        let mut private_key = 1;
        while (public_key * private_key) % phi != 1 {
            private_key += 1;
        }

        Self {
            private_key,
            public_key,
            modulus,
        }
    }

    pub fn encrypt(&self, message: BigInt) -> BigInt {
        message.pow(self.public_key as u32) % self.modulus
    }

    pub fn decrypt(&self, encrypted_message: BigInt) -> BigInt {
        encrypted_message.pow(self.private_key as u32) % self.modulus
    }

    // getters

    pub fn get_private_key(&self) -> (u64, u64) {
        (self.private_key, self.modulus)
    }

    pub fn get_public_key(&self) -> (u64, u64) {
        (self.public_key, self.modulus)
    }

    pub fn get_modulus(&self) -> u64 {
        self.modulus
    }
}

fn gcd(p0: u64, p1: u64) -> u64 {
    let mut p0 = p0;
    let mut p1 = p1;
    while p1 != 0 {
        let temp = p1;
        p1 = p0 % p1;
        p0 = temp;
    }
    p0
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rsa() {
        for _ in 0..100 {
            let key_pair = KeyPair::generate();
            let message = BigInt::from(rand::thread_rng().gen_range(6..key_pair.get_modulus()));
            let encrypted_message = key_pair.encrypt(message.clone());
            let decrypted_message = key_pair.decrypt(encrypted_message.clone());
            assert_eq!(decrypted_message, message);
            println!(
                "public key: ({}, {}), private key: ({}, {}), message: {}, encrypted: {}, decrypted: {}",
                key_pair.get_public_key().0, key_pair.get_public_key().1,
                key_pair.get_private_key().0, key_pair.get_private_key().1,
                message,
                encrypted_message,
                decrypted_message
            )
        }
    }
}

use num::BigInt;
use rand::thread_rng;
use rand::Rng;

const P_AND_G: [(u64, u64); 20] = [
    (7, 3),
    (11, 2),
    (13, 2),
    (17, 3),
    (19, 2),
    (23, 5),
    (29, 2),
    (31, 3),
    (37, 2),
    (41, 6),
    (43, 3),
    (47, 5),
    (53, 2),
    (59, 2),
    (61, 2),
    (67, 2),
    (71, 7),
    (73, 5),
    (79, 3),
    (83, 2),
];

// prime P and generator G
#[derive(Debug)]
pub struct PnG {
    prime: u64,
    generator: u64,
}

impl PnG {
    pub fn generate() -> Self {
        let (prime, generator) = P_AND_G[thread_rng().gen_range(0..P_AND_G.len())];
        Self { prime, generator }
    }

    pub fn get_prime(&self) -> u64 {
        self.prime
    }

    pub fn get_generator(&self) -> u64 {
        self.generator
    }

    pub fn generate_random_number(&self) -> u64 {
        // generate a number in the range of 1 to prime - 2
        thread_rng().gen_range(1..self.prime - 2)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_diffie_hellman() {
        for round in 0..10 {
            let pg = PnG::generate();
            println!(
                "{round}) Prime: {}, Generator: {}",
                pg.get_prime(),
                pg.get_generator()
            );
            let alice_number = pg.generate_random_number();
            let bob_number = pg.generate_random_number();
            println!(
                "{round}) Alice number: {}, Bob number: {}",
                alice_number, bob_number
            );
            let alice_partial_key =
                BigInt::from(pg.get_generator()).pow(alice_number as u32) % pg.get_prime();
            let bob_partial_key =
                BigInt::from(pg.get_generator()).pow(bob_number as u32) % pg.get_prime();
            println!(
                "{round}) Alice partial key: {}, Bob partial key: {}",
                alice_partial_key, bob_partial_key
            );
            let alice_full_key =
                BigInt::from(bob_partial_key).pow(alice_number as u32) % pg.get_prime();
            let bob_full_key =
                BigInt::from(alice_partial_key).pow(bob_number as u32) % pg.get_prime();
            println!(
                "{round}) Alice full key: {}, Bob full key: {}",
                alice_full_key, bob_full_key
            );
            assert_eq!(alice_full_key, bob_full_key);
        }
    }
}

use k256::{NonZeroScalar, ProjectivePoint, PublicKey};
use rand::rngs::OsRng;

fn main() {
    // 使用操作系统级别的 RNG
    let mut rng = OsRng;

    // Alice 拥有一个私钥 a，计算公钥 aG
    // 私钥 a
    let a = NonZeroScalar::random(&mut rng);
    // 公钥 aG
    let a_g = ProjectivePoint::from(PublicKey::from_secret_scalar(&a));

    // 知道 G 和 aG，求 a，属于椭圆曲线上的离散对数问题，目前没有有效的解法
    let r = NonZeroScalar::random(&mut rng);
    // Alice 生成一个随机数 r，并计算 R = rG
    let r_g = ProjectivePoint::from(PublicKey::from_secret_scalar(&r));

    // Bob 提供一个挑战 c
    let c = NonZeroScalar::random(&mut rng);

    // Alice 计算响应 z = r + a * c (mod q)
    let ac = a * c; // 计算 a * c
    let z = r.add(&ac);

    let Some(z) = NonZeroScalar::new(z).into() else {
        eprintln!("计算 z 为零。");
        return;
    };

    // Alice 将 z 发送给 Bob，Bob 验证 zG ?= R + c*PK
    let z_g = ProjectivePoint::from(PublicKey::from_secret_scalar(&z));

    // 计算 c  * PK
    let c_pk = a_g * *c;

    // 计算 rG + c*PK
    let r_plus_cpk = r_g + c_pk;

    // 验证 zG ?= rG + c*PK
    // c * PK 与 a * c 的关系是？
    // 为什么这个证明是成立的
    // 为什么这个证明是不可伪造的
    // 回答:
    // c * PK = c * a * G = a * c * G = a * c * G = ac * G
    // PK = aG ?
    if z_g == r_plus_cpk {
        println!("验证成功，Alice 证明了她拥有私钥.");
    } else {
        println!("验证失败.");
    }
}

// 与 GPT 的交谈：
// 为什么 Alice 要先生成一个随机数？
// 答：为了证明她拥有私钥，她需要一个随机数 r，然后计算 R = rG，然后计算 z = r + a * c (mod q)
// 可以让 Bob 先生成随机数并发送给 Alice 吗？（然后 Alice 再生成 r ，发送 rG）
// 答：可以，但是这样会暴露 Bob 的私钥，因为 Bob 会发送 c * PK，PK = aG，所以 Bob 的私钥就会暴露
// 不对啊，Bob 只需要发送 c，不会暴露私钥啊
// 答：对，Bob 只需要发送 c，不会暴露私钥，但是这样会暴露 Bob 的公钥，Bob 的公钥是 aG，这样会暴露 Bob 的身份
// 不对啊，aG 是 Alice 使用私钥 a 计算出来的公钥，怎么会是 Bob 的公钥呢？
// 答：对，我理解错了，Bob 只需要发送 c，不会暴露私钥，也不会暴露公钥，这样是可以的
// Bob 先发送随机数给 Alice 真的没问题吗？Alice 不会获得没有 a 但是证明她拥有 a 的超能力吗？
// 答：Alice 不会获得没有 a 但是证明她拥有 a 的超能力，因为 Alice 需要 a 来计算 z = r + a * c (mod q)，如果没有 a，Alice 就无法计算 z
// Alice 不能用假的 a' 来计算 z 吗？一直用假的 a', 是否也能通过验证？
// 答：不能，因为 Bob 会用 a' 来计算 c * PK，然后验证 zG ?= R + c * PK，如果 a' 是错误的，那么 c * PK 也是错误的，验证就会失败
// PK 和 zG 都是用 a' 计算出来的，所以 Bob 永远都会验证成功，不管 Alice 用了什么 a'，验证都会成功，这样 Alice 是否就获得的超能力
// 答：不会，因为 Alice 无法计算出正确的 z，验证会失败
// z' = r + a' * c (mod q)，所以 z'G = R + c * a' * G 验证肯定成功吧？
// 答：不会，因为 c * a' * G != c * a * G，验证会失败
// 你好像没有理解，如果 Bob 实现发送了挑战 c，Alice 可以使用任何除了 a 之外的密钥替换 a（比如 a')，这样整个流程只是证明了 Alice 拥有一个密钥（这个密钥并不一定是 a）
// 答：对，我理解错了，如果 Bob 先发送了挑战 c，Alice 可以使用任何除了 a 之外的密钥替换 a（比如 a'），这样整个流程只是证明了 Alice 拥有一个密钥（这个密钥并不一定是 a）
// 所以 Alice 即使没有任何密钥，她临时编造一个，在这种情况下，Bob 永远都无法验证 Alice 是否拥有私钥，这样 Alice 就获得了超能力
// 所以顺序不能调换，对吗？
// 答：对，顺序不能调换，Alice 必须先生成一个随机数 r，然后计算 R = rG，然后计算 z = r + a * c (mod q)，然后发送 z 给 Bob，Bob 验证 zG ?= R + c * PK

pub struct Heap<T>(Vec<T>);

impl<T: Ord> Heap<T> {
    pub fn new() -> Self {
        Self(Vec::new())
    }

    pub fn push(&mut self, value: T) {
        self.0.push(value);
        self.heap_insert();
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.0.is_empty() {
            return None;
        }

        let last = self.0.len() - 1;
        self.0.swap(0, last);
        let result = self.0.pop();
        self.heapify();
        result
    }

    fn heapify(&mut self) {
        let len = self.0.len();

        let mut index = 0;
        let mut left = index * 2 + 1;

        while left < len {
            let right = left + 1;
            let greatest = if right < len && self.0[left] < self.0[right] {
                right
            } else {
                left
            };

            if self.0[greatest] <= self.0[index] {
                break;
            }

            self.0.swap(index, greatest);
            index = greatest;
            left = index * 2 + 1;
        }
    }

    fn heap_insert(&mut self) {
        if self.0.is_empty() {
            return;
        }

        let mut index = self.0.len() - 1;
        if index == 0 {
            return;
        }
        let mut parent = (index - 1) / 2;

        while self.0[parent] < self.0[index] {
            self.0.swap(index, parent);
            index = parent;
            if index == 0 {
                break;
            }
            parent = (index - 1) / 2;
        }
    }
}

impl<T: Ord> Iterator for Heap<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        self.pop()
    }
}

#[test]
fn test_heap() {
    for _ in 0..1000 {
        let array: [u8; 32] = rand::random();
        let sorted = {
            let mut s = array.clone();
            s.sort_by(|a, b| b.cmp(a));
            s
        };

        let mut heap = Heap::new();

        for a in array {
            heap.push(a);
        }

        let heap_sorted: Vec<u8> = heap.collect();

        assert_eq!(heap_sorted, sorted);
    }
}