// This file is part of the laron-crypto
//
// Copyright 2023 Ade M Ramdani
//
// 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, 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 General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use k256::ecdsa;
use k256::ecdsa::signature::DigestVerifier;
use k256::ecdsa::{RecoveryId, SigningKey, VerifyingKey};
use sha3::{Digest, Keccak256};

use crate::error::{Error, Result};
use crate::{PrivateKey, PublicKey};

/// Wrapper for K256 signature.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct Signature {
    sig: [u8; 64],
    id: u8,
}

impl Signature {
    /// Create a new signature.
    pub fn new() -> Self {
        Self {
            sig: [0u8; 64],
            id: 0,
        }
    }

    /// Create a new signature from bytes.
    pub fn from_bytes(bytes: [u8; 65]) -> Self {
        let mut res = Self::new();
        res.sig.copy_from_slice(&bytes[..64]);
        res.id = bytes[64];
        res
    }

    /// Create a new signature from slice.
    pub fn from_slice(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != 64 {
            return Err(Error::InvalidSignature);
        }
        let mut res = Self::new();
        res.sig.copy_from_slice(bytes);
        if bytes.len() == 65 {
            res.id = bytes[64];
        }
        Ok(res)
    }

    /// Create a new signature from hex string.
    pub fn from_hex(hex: &str) -> Result<Self> {
        let bytes = hex::decode(hex)?;
        Self::from_slice(&bytes)
    }

    /// Get the signature as bytes.
    pub fn as_bytes(&self) -> [u8; 65] {
        let mut res = [0u8; 65];
        res[..64].copy_from_slice(&self.sig);
        res[64] = self.id;
        res
    }

    /// Get the signature as slice.
    pub fn as_slice(&self) -> &[u8] {
        let bytes = self.as_bytes();
        Box::leak(Box::new(bytes))
    }

    /// Get the signature as hex string.
    pub fn as_hex(&self) -> String {
        let bytes = self.as_bytes();
        hex::encode(bytes)
    }
}

impl From<ecdsa::Signature> for Signature {
    fn from(sig: ecdsa::Signature) -> Self {
        let mut res = Self::new();
        res.sig.copy_from_slice(sig.to_bytes().as_slice());
        res
    }
}

impl From<&Signature> for ecdsa::Signature {
    fn from(value: &Signature) -> Self {
        Self::from_slice(&value.sig).unwrap()
    }
}

impl From<&Signature> for RecoveryId {
    fn from(value: &Signature) -> Self {
        Self::from_byte(value.id).unwrap()
    }
}

/// K256Sign trait provides the interface for K256
/// Signing.
pub trait K256Sign {
    /// Sign a message with the private key
    fn sign(&self, msg: &str) -> Signature;
}

/// K256Verify trait provides the interface for K256
/// Verification.
pub trait K256Verify {
    /// Verify a signature with the public key.
    fn verify(&self, msg: &str, signature: &Signature) -> bool;
}

/// K256Recover trait provides the interface for K256
/// Recovery.
pub trait K256Recover {
    /// Recover a public key from a signature.
    fn recover(msg: &str, signature: &Signature) -> Result<PublicKey>;
}

impl From<&PrivateKey> for SigningKey {
    fn from(value: &PrivateKey) -> Self {
        Self::from_slice(value.as_slice()).unwrap()
    }
}

impl From<SigningKey> for PrivateKey {
    fn from(value: SigningKey) -> Self {
        let bytes = value.to_bytes();
        Self::from_slice(bytes.as_slice()).unwrap()
    }
}

impl From<&PublicKey> for VerifyingKey {
    fn from(value: &PublicKey) -> Self {
        Self::from_sec1_bytes(value.as_slice()).unwrap()
    }
}

impl From<VerifyingKey> for PublicKey {
    fn from(value: VerifyingKey) -> Self {
        let bytes = value.to_encoded_point(false);
        Self::from_slice(bytes.as_bytes()).unwrap()
    }
}

impl K256Sign for PrivateKey {
    fn sign(&self, msg: &str) -> Signature {
        let digest = Keccak256::new_with_prefix(msg.as_bytes());
        let key: SigningKey = self.into();
        let (sig, id) = key.sign_digest_recoverable(digest).unwrap();
        let mut signature: Signature = sig.into();
        signature.id = id.to_byte();
        signature
    }
}

impl K256Verify for PublicKey {
    fn verify(&self, msg: &str, signature: &Signature) -> bool {
        let digest = Keccak256::new_with_prefix(msg.as_bytes());
        let key: VerifyingKey = self.into();
        let sig: ecdsa::Signature = signature.into();
        key.verify_digest(digest, &sig).is_ok()
    }
}

impl K256Recover for PrivateKey {
    fn recover(msg: &str, signature: &Signature) -> Result<PublicKey> {
        let digest = Keccak256::new_with_prefix(msg.as_bytes());
        let sig: ecdsa::Signature = signature.into();
        let id: RecoveryId = signature.into();
        let recovered = VerifyingKey::recover_from_digest(digest, &sig, id)
            .map_err(|e| Error::RecoveryError(e.to_string()))?;
        let recovered_bytes = recovered.to_encoded_point(false);
        PublicKey::from_slice(recovered_bytes.as_bytes())
    }
}

impl K256Recover for PublicKey {
    fn recover(msg: &str, signature: &Signature) -> Result<PublicKey> {
        PrivateKey::recover(msg, signature)
    }
}

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

    #[test]
    fn test_k256() {
        let sk = PrivateKey::new();
        let pk = sk.public_key();
        let msg = "hello world";
        let signature = sk.sign(msg);
        assert!(pk.verify(msg, &signature));
        let recovered = PublicKey::recover(msg, &signature).unwrap();
        assert_eq!(pk, recovered);
    }
}