/* This file is part of DarkFi (https://dark.fi)
 *
 * Copyright (C) 2020-2026 Dyne.org foundation
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::convert::AsRef;

use blake2::{digest::consts::U4, Blake2b, Digest};
pub use equix::{EquiXBuilder, HashError, RuntimeOption, Solution, SolverMemory};

/// Algorithm personalization string
const P_STRING: &[u8] = b"DarkFi Equi-X\0";

/// Length of the personalization string, in bytes
const P_STRING_LEN: usize = 14;

/// Length of the nonce value generated by clients and included in the solution
pub const NONCE_LEN: usize = 16;

/// A challenge string
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Challenge(pub Vec<u8>);

impl Challenge {
    /// Build a new [`Challenge`].
    ///
    /// Copies `input` and `nonce` values into
    /// a new byte vector.
    pub fn new(input: &[u8], nonce: &[u8; NONCE_LEN]) -> Self {
        let mut result = Vec::<u8>::new();
        result.extend_from_slice(P_STRING);
        result.extend_from_slice(input.as_ref());
        result.extend_from_slice(nonce.as_ref());

        Self(result)
    }

    /// Clone the input portion of this challenge.
    pub fn input(&self) -> Vec<u8> {
        self.0[P_STRING_LEN..(self.0.len() - NONCE_LEN)].into()
    }

    /// Clone the nonce portion of this challenge.
    pub fn nonce(&self) -> [u8; NONCE_LEN] {
        self.0[(self.0.len() - NONCE_LEN)..].try_into().expect("slice length correct")
    }

    /// Increment the nonce value inside this challenge.
    pub fn increment_nonce(&mut self) {
        fn inc_le_bytes(slice: &mut [u8]) {
            for byte in slice {
                let (value, overflow) = (*byte).overflowing_add(1);
                *byte = value;
                if !overflow {
                    break;
                }
            }
        }
        let len = self.0.len();
        inc_le_bytes(&mut self.0[(len - NONCE_LEN)..]);
    }

    // Verify that a solution proof passes the effort test.
    pub fn check_effort(&self, proof: &equix::SolutionByteArray, effort: u32) -> bool {
        let mut hasher = Blake2b::<U4>::new();
        hasher.update(self.as_ref());
        hasher.update(proof.as_ref());
        let value = u32::from_be_bytes(hasher.finalize().into());
        value.checked_mul(effort).is_some()
    }
}

impl AsRef<[u8]> for Challenge {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

pub struct EquiXPow {
    /// Target effort
    pub effort: u32,
    /// The next [`Challenge`] to try
    pub challenge: Challenge,
    /// Configuration settings for Equi-X
    pub equix: EquiXBuilder,
    /// Temporary memory for Equi-X to use
    pub mem: SolverMemory,
}

impl EquiXPow {
    pub fn run(&mut self) -> Result<Solution, equix::Error> {
        loop {
            if let Some(solution) = self.run_step()? {
                return Ok(solution);
            }
        }
    }

    pub fn run_step(&mut self) -> Result<Option<Solution>, equix::Error> {
        match self.equix.build(self.challenge.as_ref()) {
            Ok(equix) => {
                for candidate in equix.solve_with_memory(&mut self.mem) {
                    if self.challenge.check_effort(&candidate.to_bytes(), self.effort) {
                        return Ok(Some(candidate))
                    }
                }
            }
            Err(equix::Error::Hash(HashError::ProgramConstraints)) => (),
            Err(e) => {
                return Err(e);
            }
        };
        self.challenge.increment_nonce();
        Ok(None)
    }

    pub fn verify(&self, challenge: &Challenge, solution: &Solution) -> Result<(), equix::Error> {
        if challenge.check_effort(&solution.to_bytes(), self.effort) {
            return self.equix.verify(challenge.as_ref(), solution)
        }
        Err(equix::Error::HashSum)
    }
}