// SPDX-FileCopyrightText: 2023 Aravinth Manivannan <realaravinth@batsense.net>
//
// SPDX-License-Identifier: Apache-2.0
// SPDX-License-Identifier: MIT

//! MCaptch's SHA256 based Proof of Work library
//!
//! # Example:
//! ```rust
//!   use mcaptcha_pow_sha256::{ConfigBuilder, PoW};
//!
//!   fn main() {
//!       let config = ConfigBuilder::default()
//!         .salt("myrandomsaltisnotlongenoug".into())
//!         .build()
//!         .unwrap();
//!
//!       let phrase = "ironmansucks";
//!
//!       const DIFFICULTY: u32 = 1000;
//!
//!       let work = config.prove_work(&phrase, DIFFICULTY).unwrap();
//!       assert!(config.is_valid_proof(&work, &phrase));
//!       assert!(config.is_sufficient_difficulty(&work, DIFFICULTY));
//!   }    
//! ```

use std::marker::PhantomData;

use derive_builder::Builder;
//use num::Num;
use serde::{Deserialize, Serialize};
use serde_wincode::SerdeCompat;
use sha2::{digest::Update, Digest, Sha256};

fn serialize_input<T: Serialize>(t: &T) -> Result<Vec<u8>, wincode::error::WriteError> {
    <SerdeCompat<T> as wincode::Serialize>::serialize(t)
}

/// Proof of Work over concrete type T. T can be any type that implements serde::Serialize.
#[derive(Serialize, Builder, Deserialize, PartialEq, Clone, Debug)]
pub struct PoW<T> {
    pub nonce: u64,
    pub result: String,
    #[builder(default = "PhantomData", setter(skip))]
    _spook: PhantomData<T>,
}

/// Configuration for generting proof of work
/// Please choose a long, unique value for salt
/// Resistance to dictionary/rainbow attacks depend on uniqueness
/// of the salt
#[derive(Serialize, Deserialize, Builder, PartialEq, Clone, Debug)]
pub struct Config {
    pub salt: String,
}

/// Return value of incremental prooving. When proof  is ready, IncrementalSolve::Work is returned
/// and when proof is not ready, IncrementalSolve::Intermediate is returned
#[cfg(feature = "incremental")]
pub enum IncrementalSolve<T> {
    /// Intermediate result
    Intermediate(u128, u64, Sha256, u128),
    /// Final result
    Work(PoW<T>),
}

impl Config {
    ///
    /// step is used to control the number of cycles after which the function should exit, even
    /// when the proof isn't ready
    /// inter is used to keep track of state and complete proof generation. Set inter to None
    /// during first cyle and pass the returned value of the previous cycle to continue proof
    /// generation

    #[cfg(feature = "incremental")]
    pub fn stepped_prove_work<T>(
        &self,
        t: &T,
        difficulty: u32,
        step: usize,
        inter: Option<IncrementalSolve<T>>,
    ) -> Result<IncrementalSolve<T>, wincode::error::WriteError>
    where
        T: Serialize,
    {
        let v = serialize_input(t)?;
        Ok(self.stepped_prove_work_serialized(&v, difficulty, step, inter))
    }

    #[cfg(feature = "incremental")]
    /// Create Proof of Work over item of type T.
    ///
    /// Make sure difficulty is not too high. A 64 bit difficulty,
    /// for example, takes a long time on a general purpose processor.
    /// The input is assumed to be serialized using network byte order.
    ///
    /// Make sure difficulty is not too high. A 64 bit difficulty,
    /// for example, takes a long time on a general purpose processor.
    /// step is used to control the number of cycles after which the function should exit, even
    /// when the proof isn't ready
    /// inter is used to keep track of state and complete proof generation. Set inter to None
    /// during first cyle and pass the returned value of the previous cycle to continue proof
    /// generation
    /// Returns an error if serialization fails.
    /// Create Proof of Work on an already serialized item of type T.
    fn stepped_prove_work_serialized<T>(
        &self,
        prefix: &[u8],
        difficulty: u32,
        step: usize,
        inter: Option<IncrementalSolve<T>>,
    ) -> IncrementalSolve<T>
    where
        T: Serialize,
    {
        let (mut result, mut n, prefix_sha, difficulty) = match inter {
            Some(IncrementalSolve::Intermediate(result, nonce, prefix, difficulty)) => {
                (result, nonce, prefix, difficulty)
            }
            _ => {
                let prefix_sha = Sha256::new().chain(&self.salt).chain(prefix);
                let n = 0;
                let result = 0;
                let difficulty = get_difficulty(difficulty);
                (result, n, prefix_sha, difficulty)
            }
        };
        let mut count = 0;
        while result < difficulty {
            if count > step {
                return IncrementalSolve::Intermediate(result, n, prefix_sha, difficulty);
            } else {
                count += 1;
            }
            n += 1;
            result = dev::score(prefix_sha.clone(), n);
        }
        IncrementalSolve::Work(PoW {
            nonce: n,
            result: result.to_string(),
            _spook: PhantomData,
        })
    }

    /// Create Proof of Work over item of type T.
    ///
    /// Make sure difficulty is not too high. A 64 bit difficulty,
    /// for example, takes a long time on a general purpose processor.
    /// Returns an error if serialization fails.
    pub fn prove_work<T>(
        &self,
        t: &T,
        difficulty: u32,
    ) -> Result<PoW<T>, wincode::error::WriteError>
    where
        T: Serialize,
    {
        serialize_input(t).map(|v| self.prove_work_serialized(&v, difficulty))
    }

    /// Create Proof of Work on an already serialized item of type T.
    /// The input is assumed to be serialized using network byte order.
    ///
    /// Make sure difficulty is not too high. A 64 bit difficulty,
    /// for example, takes a long time on a general purpose processor.
    pub fn prove_work_serialized<T>(&self, prefix: &[u8], difficulty: u32) -> PoW<T>
    where
        T: Serialize,
    {
        let prefix_sha = Sha256::new().chain(&self.salt).chain(prefix);
        let mut n = 0;
        let mut result = 0;
        let difficulty = get_difficulty(difficulty);
        while result < difficulty {
            n += 1;
            result = dev::score(prefix_sha.clone(), n);
        }
        PoW {
            nonce: n,
            result: result.to_string(),
            _spook: PhantomData,
        }
    }

    /// Calculate the PoW score with the provided input T.
    pub fn calculate<T>(
        &self,
        pow: &PoW<T>,
        t: &T,
    ) -> Result<u128, wincode::error::WriteError>
    where
        T: Serialize,
    {
        serialize_input(t).map(|v| self.calculate_serialized(pow, &v))
    }

    /// Calculate the PoW score of an already serialized T and self.
    /// The input is assumed to be serialized using network byte order.
    pub fn calculate_serialized<T>(&self, pow: &PoW<T>, target: &[u8]) -> u128
    where
        T: Serialize,
    {
        dev::score(Sha256::new().chain(&self.salt).chain(target), pow.nonce)
    }

    /// Verifies that the PoW is indeed generated out of the phrase provided.
    pub fn is_valid_proof<T>(&self, pow: &PoW<T>, t: &T) -> bool
    where
        T: Serialize,
    {
        match self.calculate(pow, t) {
            Ok(res) => {
                return if pow.result == res.to_string() {
                    true
                } else {
                    false
                }
            }
            Err(_) => return false,
        }
    }

    /// Checks if the PoW result is of sufficient difficulty
    pub fn is_sufficient_difficulty<T>(&self, pow: &PoW<T>, target_diff: u32) -> bool
    where
        T: Serialize,
    {
        match pow.result.parse::<u128>() {
            Ok(res) => return res >= get_difficulty(target_diff),
            Err(_) => return false,
        }
    }
}

pub mod dev {
    use super::*;

    pub fn score(prefix_sha: Sha256, nonce: u64) -> u128 {
        first_bytes_as_u128(
            prefix_sha
                .chain(&nonce.to_string()) //used to be: to_be_bytes() converts to network endian
                // chain() expexets something that can be converted to &[u8], String is fine
                .finalize()
                .as_slice(),
        )
    }

    /// # Panics
    ///
    /// panics if inp.len() < 16
    fn first_bytes_as_u128(inp: &[u8]) -> u128 {
        let mut bytes = [0u8; 16];
        bytes.copy_from_slice(&inp[..16]);
        u128::from_be_bytes(bytes)
    }
}

// utility function to get u128 difficulty factor from u32
// javacript isn't capable of represnting u128 so
fn get_difficulty(difficulty_factor: u32) -> u128 {
    u128::MAX - u128::MAX / difficulty_factor as u128
}

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

    const DIFFICULTY: u32 = 1000;

    fn get_config() -> Config {
        ConfigBuilder::default()
            .salt(
                "79ziepia7vhjgviiwjhnend3ofjqocsi2winc4ptqhmkvcajihywxcizewvckg9h6gs4j83v9".into(),
            )
            .build()
            .unwrap()
    }

    #[test]
    fn base_functionality() {
        // Let's prove we did work targeting a phrase.
        let phrase = b"Ex nihilo nihil fit.".to_vec();
        let config = get_config();
        let pw = config.prove_work(&phrase, DIFFICULTY).unwrap();
        assert!(config.calculate(&pw, &phrase).unwrap() >= get_difficulty(DIFFICULTY));
        assert!(config.is_valid_proof(&pw, &phrase));
        assert!(config.is_sufficient_difficulty(&pw, DIFFICULTY));
    }

    #[test]
    fn double_pow() {
        let phrase = "Ex nihilo nihil fit.".to_owned();
        let config = get_config();

        let pw = config.prove_work(&phrase, DIFFICULTY).unwrap();
        let pwpw = config.prove_work(&pw, DIFFICULTY).unwrap();

        assert!(config.calculate(&pw, &phrase).unwrap() >= get_difficulty(DIFFICULTY));
        assert!(config.is_valid_proof(&pw, &phrase));
        assert!(config.is_sufficient_difficulty(&pw, DIFFICULTY));

        assert!(config.calculate(&pwpw, &pw).unwrap() >= get_difficulty(DIFFICULTY));
        assert!(config.is_valid_proof(&pwpw, &pw));
        assert!(config.is_sufficient_difficulty(&pwpw, DIFFICULTY));
    }

    #[test]
    fn is_not_valid_proof() {
        let phrase = "Ex nihilo nihil fit.".to_owned();
        let phrase2 = "Omne quod movetur ab alio movetur.".to_owned();

        let config = get_config();
        let pw = config.prove_work(&phrase, DIFFICULTY).unwrap();

        let pw2 = config.prove_work(&phrase2, DIFFICULTY).unwrap();

        assert!(!config.is_valid_proof(&pw, &phrase2));
        assert!(!config.is_valid_proof(&pw2, &phrase));
    }

    #[test]
    fn difficulty_monotonicity() {
        const D_LOW: u32 = 1000;
        const D_HIGH: u32 = 10_000;
        let target = "testing";
        let config = get_config();
        let pw_low = config.prove_work(&target, D_LOW).unwrap();
        let pw_high = config.prove_work(&target, D_HIGH).unwrap();
        assert!(config.is_sufficient_difficulty(&pw_low, D_LOW));
        assert!(config.is_sufficient_difficulty(&pw_high, D_HIGH));
        let score_low = config.calculate(&pw_low, &target).unwrap();
        let score_high = config.calculate(&pw_high, &target).unwrap();
        assert!(score_low >= get_difficulty(D_LOW));
        assert!(score_high >= get_difficulty(D_HIGH));
        assert!(get_difficulty(D_HIGH) > get_difficulty(D_LOW));
    }

    #[test]
    fn serialization_test() {
        let target: u8 = 1;
        let config = get_config();
        let pw = config.prove_work(&target, DIFFICULTY).unwrap();

        let message: (u8, PoW<u8>) = (target, pw);
        let message_ser = serialize_input(&message).unwrap();
        let recieved_message =
            <SerdeCompat<(u8, PoW<u8>)> as wincode::Deserialize>::deserialize(&message_ser)
                .unwrap();
        assert_eq!(recieved_message, message);
        assert!(config.is_sufficient_difficulty(&message.1, DIFFICULTY));
        assert!(config.is_valid_proof(&message.1, &target));
    }

    #[test]
    #[cfg(feature = "incremental")]
    fn stepped_solve() {
        let phrase = "Ex nihilo nihil fit.".to_owned();
        let config = get_config();

        let mut inter = None;
        loop {
            match config.stepped_prove_work(&phrase, 50000, 1000, inter) {
                Ok(IncrementalSolve::Intermediate(result, nonce, prefix, difficulty)) => {
                    println!("Current nonce {nonce}");
                    inter = Some(IncrementalSolve::Intermediate(
                        result, nonce, prefix, difficulty,
                    ));
                    continue;
                }

                Ok(IncrementalSolve::Work(w)) => {
                    assert!(config.is_valid_proof(&w, &phrase));
                    assert!(config.is_sufficient_difficulty(&w, DIFFICULTY));
                    break;
                }
                Err(e) => panic!("{}", e),
            };
        }
    }
}
