//! Just The Math Crypto Library
//!
// Copyright (C) 2020 Michael Mestnik
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; version 2 of the License
//
// 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 General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
//
//! Don't use/insecure.  If you must, use it for a pet project or to validate another implementation.

pub mod brainpool_p256r1 {
    //! https://andrea.corbellini.name/2015/05/23/elliptic-curve-cryptography-finite-fields-and-discrete-logarithms
    use num_bigint_dig::{BigInt, BigUint, Sign::*};
    use std::ops::{Add, AddAssign, Mul};

    pub mod util {
        use num_bigint_dig::{BigInt, BigUint, Sign::Plus};
        pub fn egcd(a: BigInt, b: BigInt) -> (BigInt, BigInt, BigInt) {
            if a == 0.into() {
                (b, 0.into(), 1.into())
            } else {
                use num_integer::Integer;
                let (g, y, x) = egcd(b.mod_floor(&a), a.clone());
                (g, x - (b / a) * y.clone(), y)
            }
        }
        pub fn mod_inv(a: BigInt, p: BigUint) -> BigUint {
            if a < 0.into() {
                p.clone() - mod_inv(-a, p.clone())
            } else {
                let p = BigInt::from_biguint(Plus, p);
                let (g, x, _y) = egcd(a, p.clone());
                if g == 1.into() {
                    use num_integer::Integer;
                    x.mod_floor(&p).to_biguint().unwrap()
                } else {
                    panic!()
                }
            }
        }
        // https://reinerm.wordpress.com/programming/rust/number-theory-in-rust/

        // Note that % can yield negative results.
        pub fn modulo(n: BigInt, m: BigInt) -> BigInt {
            use num_integer::Integer;
            n.mod_floor(&m)
        }
        // Compute a^n mod m.
        pub fn power_mod(base: BigInt, exp: BigInt, modulus: BigInt) -> BigInt {
            use num_integer::Integer;
            if exp < 0.into() {
                unimplemented!()
            }
            let mut b: BigInt = base.mod_floor(&modulus);
            let mut result: BigInt = 1.into();
            let mut e: BigInt = exp;
            while e > 0.into() {
                if e.is_odd() {
                    result = (result * b.clone()).mod_floor(&modulus);
                }
                b = (b.clone() * b.clone()).mod_floor(&modulus);
                e >>= 1;
            }
            modulo(result, modulus)
        }
        // Legendre symbol, returns 1, 0, or -1 mod p
        pub fn ls(a: BigInt, p: BigInt) -> BigInt {
            power_mod(a, (p.clone() - 1) / 2, p)
        }
    }

    pub mod paramiters {
        //! https://tools.ietf.org/html/rfc5639#section-3.4
        pub const P: &[u8] = b"A9FB57DBA1EEA9BC3E660A909D838D726E3BF623D52620282013481D1F6E5377";
        pub const A: &[u8] = b"7D5A0975FC2C3057EEF67530417AFFE7FB8055C126DC5C6CE94A4B44F330B5D9";
        pub const B: &[u8] = b"26DC5C6CE94A4B44F330B5D9BBD77CBF958416295CF7E1CE6BCCDC18FF8C07B6";
        pub const X: &[u8] = b"8BD2AEB9CB7E57CB2C4B482FFC81B7AFB9DE27E1E3BD23C23A4453BD9ACE3262";
        pub const Y: &[u8] = b"547EF835C3DAC4FD97F8461A14611DC9C27745132DED8E545C1D54C72F046997";
        pub const Q: &[u8] = b"A9FB57DBA1EEA9BC3E660A909D838D718C397AA3B561A6F7901E0E82974856A7";
        pub const H: u8 = 1;
    }

    pub fn get_p() -> BigUint {
        BigUint::parse_bytes(paramiters::P, 16).unwrap()
    }
    pub fn get_int_p() -> BigInt {
        BigInt::from_biguint(Plus, BigUint::parse_bytes(paramiters::P, 16).unwrap())
    }
    pub fn mod_floor_p(m: BigUint) -> BigUint {
        use num_integer::Integer;
        m.mod_floor(&BigUint::parse_bytes(paramiters::P, 16).unwrap())
    }
    pub fn mod_floor_int_p(m: BigInt) -> BigInt {
        use num_integer::Integer;
        m.mod_floor(&BigInt::from_biguint(
            Plus,
            BigUint::parse_bytes(paramiters::P, 16).unwrap(),
        ))
    }
    pub fn get_a() -> BigUint {
        BigUint::parse_bytes(paramiters::A, 16).unwrap()
    }
    pub fn get_int_a() -> BigInt {
        BigInt::from_biguint(Plus, BigUint::parse_bytes(paramiters::A, 16).unwrap())
    }
    pub fn get_b() -> BigInt {
        BigInt::from_biguint(Plus, BigUint::parse_bytes(paramiters::B, 16).unwrap())
    }
    pub fn get_base_x() -> BigInt {
        BigInt::from_biguint(Plus, BigUint::parse_bytes(paramiters::X, 16).unwrap())
    }
    pub fn get_base_y() -> BigInt {
        BigInt::from_biguint(Plus, BigUint::parse_bytes(paramiters::Y, 16).unwrap())
    }
    pub fn get_g() -> Point {
        Point::P(
            BigUint::parse_bytes(paramiters::X, 16).unwrap(),
            BigUint::parse_bytes(paramiters::Y, 16).unwrap(),
        )
    }
    pub fn get_q() -> BigUint {
        BigUint::parse_bytes(paramiters::Q, 16).unwrap()
    }
    pub fn get_int_q() -> BigInt {
        BigInt::from_biguint(Plus, BigUint::parse_bytes(paramiters::Q, 16).unwrap())
    }
    pub fn mod_floor_q(m: BigUint) -> BigUint {
        use num_integer::Integer;
        m.mod_floor(&BigUint::parse_bytes(paramiters::Q, 16).unwrap())
    }
    pub fn mod_floor_int_q(m: BigInt) -> BigUint {
        use num_integer::Integer;
        m.mod_floor(&BigInt::from_biguint(
            Plus,
            BigUint::parse_bytes(paramiters::Q, 16).unwrap(),
        ))
        .to_biguint()
        .unwrap()
    }
    /// https://github.com/alexmgr/tinyec/blob/a812431e7881ce9b509e5fe6322b401ce3235299/tinyec/ec.py#L37
    pub fn on_curve(x: &BigInt, y: &BigInt) -> bool {
        use num_traits::{Pow, Zero};
        BigInt::zero() == mod_floor_int_p(y.pow(2u32) - x.pow(3u32) - (get_int_a() * x) - get_b())
    }

    #[derive(Clone, PartialEq)]
    pub enum Point {
        P(BigUint, BigUint),
        Inf,
    }
    use std::fmt;
    impl fmt::Debug for Point {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                Point::P(x, y) => {
                    write!(f, "Point::P({}, {}) on_curve = {}", x, y, self.on_curve())
                }
                Point::Inf => write!(f, "Point::Inf"),
            }
        }
    }
    impl Point {
        pub fn on_curve(&self) -> bool {
            if let Self::P(x, y) = self {
                on_curve(
                    &BigInt::from_biguint(Plus, x.clone()),
                    &BigInt::from_biguint(Plus, y.clone()),
                )
            } else {
                true // unimplemented!()
            }
        }
    }
    /// https://github.com/alexmgr/tinyec/blob/a812431e7881ce9b509e5fe6322b401ce3235299/tinyec/ec.py#L126
    impl Add for Point {
        type Output = Point;

        fn add(self, rhs: Self) -> Self::Output {
            assert!(self.on_curve());
            if self == rhs {
                return self.double();
            }
            assert!(rhs.on_curve());
            if let Point::P(x1, y1) = self.clone() {
                if let Point::P(x2, y2) = rhs {
                    use num_traits::Pow;
                    if x1 == x2 && y1 != y2 {
                        return Point::Inf;
                    }
                    let slope = (BigInt::from_biguint(Plus, y1.clone())
                        + BigInt::from_biguint(Minus, y2.clone()))
                        * BigInt::from_biguint(
                            Plus,
                            util::mod_inv(
                                BigInt::from_biguint(Plus, x1.clone())
                                    + BigInt::from_biguint(Minus, x2.clone()),
                                get_p(),
                            ),
                        );
                    let ret_x = mod_floor_int_p(
                        slope.pow(2u32)
                            + BigInt::from_biguint(Minus, x1.clone())
                            + BigInt::from_biguint(Minus, x2),
                    );
                    Point::P(
                        ret_x.clone().to_biguint().unwrap(),
                        mod_floor_int_p(
                            -(BigInt::from_biguint(Plus, y1)
                                + slope * (ret_x + BigInt::from_biguint(Minus, x1))),
                        )
                        .to_biguint()
                        .unwrap(),
                    )
                } else {
                    self
                }
            } else {
                rhs
            }
        }
    }
    /// https://github.com/alexmgr/tinyec/blob/a812431e7881ce9b509e5fe6322b401ce3235299/tinyec/ec.py#L124
    impl Point {
        pub fn double(&self) -> Self {
            if let Point::P(x, y) = self.clone() {
                use num_traits::Pow;
                let slope = (3u32 * x.pow(2u32) + get_a())
                    * util::mod_inv(BigInt::from_biguint(Plus, 2u32 * y.clone()), get_p());
                let ret_x = mod_floor_p(slope.pow(2u32) - 2u32 * x.clone());
                Point::P(
                    ret_x.clone(),
                    mod_floor_int_p(
                        -(BigInt::from_biguint(Plus, y)
                            + BigInt::from_biguint(Plus, slope)
                                * (BigInt::from_biguint(Plus, ret_x)
                                    + BigInt::from_biguint(Minus, x))),
                    )
                    .to_biguint()
                    .unwrap(),
                )
            } else {
                Point::Inf
            }
        }
    }
    impl AddAssign<Point> for Point {
        fn add_assign(&mut self, rhs: Point) {
            *self = self.clone() + rhs;
        }
    }
    impl Mul<BigUint> for Point {
        type Output = Point;
        fn mul(self, rhs: BigUint) -> Self::Output {
            self * BigInt::from_biguint(Plus, rhs)
        }
    }
    impl Mul<BigInt> for Point {
        type Output = Point;
        /// https://github.com/alexmgr/tinyec/blob/a812431e7881ce9b509e5fe6322b401ce3235299/tinyec/ec.py#L162
        fn mul(self, rhs: BigInt) -> Self::Output {
            if mod_floor_int_q(rhs.clone()) != 0u32.into() {
                if let Point::P(x, y) = &self {
                    use num_traits::Zero;
                    assert!(self.on_curve());
                    let mut addend = if rhs < BigInt::zero() {
                        Point::P(
                            x.clone(),
                            mod_floor_int_p(BigInt::from_biguint(Minus, y.clone()))
                                .to_biguint()
                                .unwrap(),
                        )
                    } else {
                        self
                    };
                    let mut result = Point::Inf;
                    rhs.to_radix_be(2).1.iter().rev().for_each(|b| {
                        if *b == 1 {
                            result += addend.clone()
                        }
                        addend += addend.clone()
                    });
                    result
                } else {
                    Point::Inf
                }
            } else {
                Point::Inf
            }
        }
    }
    impl From<BigInt> for Point {
        fn from(d: BigInt) -> Self {
            get_g() * d
        }
    }
    impl From<BigUint> for Point {
        fn from(d: BigUint) -> Self {
            get_g() * d
        }
    }
    impl Into<[u8; 33]> for Point {
        fn into(self) -> [u8; 33] {
            if let Point::P(x, y) = &self {
                use num_integer::Integer;
                let parity = y.is_even();
                let mut data = x.to_bytes_le();
                data.resize_with(32, Default::default);
                data.push(if parity { 2u8 } else { 3u8 });
                let data: Vec<u8> = data.into_iter().rev().collect();
                let mut array = [0; 33];
                let data = &data[..array.len()];
                array.copy_from_slice(data);
                array
            } else {
                unimplemented!()
            }
        }
    }
    impl From<[u8; 33]> for Point {
        fn from(data: [u8; 33]) -> Self {
            use num_bigint_dig::IntoBigInt;
            use num_integer::Integer;
            use num_traits::pow::Pow;

            let mut iter = data.to_vec().into_iter();
            let magic = iter.next().unwrap();
            let parity1 = match magic {
                2u8 => true,
                3u8 => false,
                _ => unimplemented!(),
            };
            let x = BigUint::from_radix_be(&iter.collect::<Vec<u8>>(), 256).unwrap();
            let y = x.pow(3u32) + (get_a() * x.clone()) + get_b().to_biguint().unwrap();
            let y = util::power_mod(y.into_bigint().unwrap(), (get_int_p() + 1) / 4, get_int_p())
                .to_biguint()
                .unwrap();
            let parity2 = y.is_even();
            let y = if (parity1 && parity2) || (!parity1 && !parity2) {
                y
            } else {
                get_p() - y
            };
            Point::P(x, y)
        }
    }

    pub fn ecdsa_sign<F>(z: BigUint, d: BigUint, mut f: F) -> (BigUint, BigUint)
    where
        F: FnMut(BigUint) -> BigUint,
    {
        let z = mod_floor_p(z);
        loop {
            use num_traits::Zero;
            let k = f(get_q() - 1u32);
            assert!(k < get_q() - 1u32);
            if let Point::P(px, _) = get_g() * k.clone() {
                use num_bigint_dig::traits::ModInverse;
                let r = mod_floor_q(px);
                if r == BigUint::zero() {
                    continue;
                }
                let s = mod_floor_q(
                    k.clone()
                        .mod_inverse(&get_q())
                        .unwrap()
                        .to_biguint()
                        .unwrap()
                        * mod_floor_q(z.clone() + mod_floor_q(r.clone() * d.clone())),
                );
                if s == BigUint::zero() {
                    continue;
                }
                return (r, s);
            } else {
                unreachable!()
            }
        }
    }
    pub fn ecdsa_verify(h: Point, z: BigUint, r: BigUint, s: BigUint) -> bool {
        let z = mod_floor_p(z);
        use num_bigint_dig::traits::ModInverse;
        use num_traits::Zero;
        assert!(h.on_curve());
        let sinv = s.mod_inverse(&get_q()).unwrap().to_biguint().unwrap();
        let u1 = mod_floor_q(sinv.clone() * z);
        let u2 = mod_floor_q(sinv * r.clone());
        assert_ne!(u1, BigUint::zero());
        assert_ne!(u2, BigUint::zero());
        if let Point::P(px, _) = (get_g() * u1) + (h * u2) {
            r == mod_floor_q(px)
        } else {
            false
        }
    }

    #[cfg(test)]
    mod tests {
        use crate::brainpool_p256r1 as bp;
        use num_bigint_dig::BigUint;

        #[test]
        #[ignore]
        fn debug() {}

        #[test]
        fn double() -> Result<(), ()> {
            double_impl().map(|_| ())
        }

        fn double_impl() -> Result<(bp::Point, bp::Point), ()> {
            let pp = bp::get_g().double();
            assert_eq!(
                pp,
                bp::Point::P(
                    BigUint::from_radix_be(
                        &base64::decode("dDzxuLXNTy61X4qjaVk6xDbvBEFmaZ431RoUws4T6g4=").unwrap(),
                        256
                    )
                    .unwrap(),
                    BigUint::from_radix_be(
                        &base64::decode("Nu0WMzfeupyUb+C7d2Up2jjfBZ9pJJQGiSraCX7rfNQ=").unwrap(),
                        256
                    )
                    .unwrap()
                )
            );
            let pppp = pp.double();
            assert_eq!(
                pppp,
                bp::Point::P(
                    BigUint::from_radix_be(
                        &base64::decode("NnIDC6znh6oxniHUBkWymZAGvuxDf9CE3T/FkvX813w=").unwrap(),
                        256
                    )
                    .unwrap(),
                    BigUint::from_radix_be(
                        &base64::decode("M1sibOX6wMNqGM5C6V9Dye7T4la90MmOVaBpWVUV0Vs=").unwrap(),
                        256
                    )
                    .unwrap()
                )
            );
            if pp.on_curve() && pppp.on_curve() {
                Ok((pp, pppp))
            } else {
                Err(())
            }
        }

        #[test]
        fn add() -> Result<(), ()> {
            add_impl().map(|_| ())
        }

        fn add_impl() -> Result<bp::Point, ()> {
            if let Ok((pp, pppp)) = double_impl() {
                assert_eq!(
                    bp::get_g() + pp.clone(),
                    bp::Point::P(
                        BigUint::from_radix_be(
                            &base64::decode("qPIXt3M48dTWYkw6tPbMFtKqhD0MD8oBa5HirSXK450=")
                                .unwrap(),
                            256
                        )
                        .unwrap(),
                        BigUint::from_radix_be(
                            &base64::decode("S0nK/H2sJrsKoqaFChtA9frBDkWJNI+3fmXMVgK3T50=")
                                .unwrap(),
                            256
                        )
                        .unwrap()
                    )
                );
                assert_eq!(
                    pp.clone() + bp::get_g(),
                    bp::Point::P(
                        BigUint::from_radix_be(
                            &base64::decode("qPIXt3M48dTWYkw6tPbMFtKqhD0MD8oBa5HirSXK450=")
                                .unwrap(),
                            256
                        )
                        .unwrap(),
                        BigUint::from_radix_be(
                            &base64::decode("S0nK/H2sJrsKoqaFChtA9frBDkWJNI+3fmXMVgK3T50=")
                                .unwrap(),
                            256
                        )
                        .unwrap()
                    )
                );
                let r1 = pp + pppp;
                assert!(r1.on_curve());
                let r2 = (bp::get_g() + bp::get_g().double()).double();
                assert_eq!(r1, r2);
                if r1.on_curve() && r1 == r2 {
                    Ok(r1)
                } else {
                    Err(())
                }
            } else {
                panic!("Requiers double()");
            }
        }

        #[test]
        fn multiply() {
            if let Ok(r) = add_impl() {
                let p: bp::Point = bp::get_g() * BigUint::from(0b110u32);
                assert_eq!(p, r);
            } else {
                panic!("Requiers add()");
            }
        }

        #[test]
        fn paramiters() {
            use num_traits::One;
            let g1 = bp::get_g();
            assert!(g1.on_curve());
            assert_eq!(bp::get_g() * bp::get_q(), bp::Point::Inf);
            let g2 = bp::get_g() * (bp::get_q() + BigUint::one());
            assert_eq!(g2, g1);
        }

        #[test]
        fn point_serialization() {
            let gorig = bp::get_g();
            let gserial: [u8; 33] = gorig.clone().into();
            assert_eq!(
                &gserial[..],
                &[
                    3, 139, 210, 174, 185, 203, 126, 87, 203, 44, 75, 72, 47, 252, 129, 183, 175,
                    185, 222, 39, 225, 227, 189, 35, 194, 58, 68, 83, 189, 154, 206, 50, 98
                ][..]
            );
            let gnew: super::Point = gserial.into();
            assert_eq!(gnew, gorig);
        }

        #[test]
        fn ecdsa() {
            let (r, s) = bp::ecdsa_sign(3u32.into(), 4u32.into(), |_| 5u32.into());
            assert!(bp::ecdsa_verify(
                bp::get_g() * BigUint::from(4u32),
                3u32.into(),
                r,
                s
            ));
        }

        #[test]
        #[ignore]
        fn conversion() {
            /*            println!(
                "{}",
                base64::encode(&BigUint::parse_bytes(b"24626335408888534556217428905503724368080122560765218856066532503467218229116", 10).unwrap().to_radix_be(256))
            ); */
            println!(
                "{}",
                base64::encode(
                    &BigUint::parse_bytes(
                        b"10000000000000000000000000000000000000000000000000000000000000000",
                        16
                    )
                    .unwrap()
                    .to_radix_be(256)
                )
            );
            /*            println!(
                "{}",
                BigUint::from_radix_be(
                    &base64::decode("NnIDC6znh6oxniHUBkWymZAGvuxDf9CE3T/FkvX813w=").unwrap(),
                    256
                )
                .unwrap()
            );
            println!(
                "{}",
                BigUint::from_radix_be(
                    &base64::decode("VP2r7dD3VN4fMwVITsHGuTcd+xHqkxAUEAmkDo+3Kbs=").unwrap(),
                    256
                )
                .unwrap()
            ); */
        }
    }
}

pub mod keys {
    pub use crate::brainpool_p256r1 as curve;
    pub use curve::Point;
    use num_bigint_dig::BigUint;
    type HmacSha512 = hmac::Hmac<sha2::Sha512>;
    type ChainCode = [u8; 32];

    /// bip-0032 uses a method that's bias to low numbers.
    /// This equation distributes the set evenly.
    pub fn get_adjustment(d: [u8; 32]) -> BigUint {
        use num_integer::Integer;
        let modulo = curve::get_q() - 1u32;
        ((BigUint::from_radix_be(&d, 256).unwrap() * modulo.clone())
            / BigUint::parse_bytes(
                b"10000000000000000000000000000000000000000000000000000000000000000",
                16,
            )
            .unwrap())
        .mod_floor(&modulo)
            + 1u32
    }

    pub fn pub_advance(p: curve::Point, mut mac: HmacSha512, i: u32) -> (BigUint, [u8; 32]) {
        if let curve::Point::P(x, y) = p {
            use hmac::Mac;
            use std::convert::TryInto;
            let mut x_data = x.to_bytes_le();
            x_data.resize_with(32, Default::default);
            let mut y_data = y.to_bytes_le();
            y_data.resize_with(32, Default::default);
            let mut data: Vec<u8> = y_data.into_iter().chain(x_data.into_iter()).rev().collect();
            data.extend_from_slice(&i.to_be_bytes());
            mac.input(&data);
            let result = mac.result().code();
            let result = result.as_slice();
            (
                get_adjustment(result[0..32].try_into().unwrap()),
                result[32..64].try_into().unwrap(),
            )
        } else {
            unimplemented!()
        }
    }

    #[derive(Clone, Debug)]
    pub struct PrivKey {
        pub d: BigUint,
        pub chain_code: ChainCode,
        h: Option<Point>,
    }

    impl PrivKey {
        pub fn new(d: BigUint, chain_code: ChainCode) -> Self {
            Self {
                d,
                chain_code,
                h: None,
            }
        }

        pub fn get_h(&mut self) -> &mut Point {
            let d = self.d.clone();
            self.h.get_or_insert_with(move || Point::from(d))
        }
    }

    impl PrivKey {
        pub fn ckd_priv(mut self, i: i32) -> Self {
            use curve::mod_floor_q;
            use hmac::Mac;
            use std::convert::TryInto;
            let mut mac = HmacSha512::new_varkey(&self.chain_code).unwrap();
            if i < 0 {
                let i = -i;
                let mut data = self.d.to_bytes_le();
                data.resize_with(33, Default::default);
                let mut data: Vec<u8> = data.into_iter().rev().collect();
                data.extend_from_slice(&i.to_be_bytes());
                mac.input(&data);
                let result = mac.result().code();
                let result = result.as_slice();

                Self::new(
                    mod_floor_q(self.d + get_adjustment(result[0..32].try_into().unwrap())),
                    result[32..64].try_into().unwrap(),
                )
            } else {
                let (d, chain_code) = pub_advance(self.get_h().clone(), mac, i as _);
                Self::new(mod_floor_q(self.d + d), chain_code)
            }
        }
    }

    #[derive(Clone, Debug, PartialEq)]
    pub struct PubKey(pub Point, pub ChainCode);

    impl Into<PubKey> for PrivKey {
        fn into(mut self) -> PubKey {
            let chain_code = self.chain_code;
            PubKey(self.get_h().clone(), chain_code)
        }
    }

    impl PubKey {
        pub fn ckd_pub(self, i: u32) -> Self {
            use hmac::Mac;
            let mac = HmacSha512::new_varkey(&self.1).unwrap();
            let (d, chain_code) = pub_advance(self.0.clone(), mac, i);
            Self(self.0 + d.into(), chain_code)
        }
    }

    #[cfg(test)]
    mod test {
        #[test]
        fn try_this() {
            let base = super::PrivKey::new(1u32.into(), [1; 32]);
            let pub_key: super::PubKey = base.clone().into();
            assert_eq!(pub_key.ckd_pub(1), base.ckd_priv(1).into());
        }
    }
}