use std::{borrow::Cow, io::Write, path::Path};

use ed25519_dalek::{ed25519::signature::SignerMut, pkcs8::DecodePrivateKey};
use sha2::Digest;

use super::V1_SIG_ED25519;

pub struct SignEd25519 {
    signing_key: ed25519_dalek::SigningKey,
    verifying_key: ed25519_dalek::VerifyingKey,
}

pub struct VerifyEd25519 {
    verifying_key: ed25519_dalek::VerifyingKey,
}

pub enum SignatureMethod {
    Ed25519(SignEd25519),
}

pub enum VerificationMethod {
    Ed25519(VerifyEd25519),
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("ed25519 signature error: {0}")]
    Ed25519(#[from] ed25519_dalek::SignatureError),
    #[error("Cannot fit signature into the message: buffer size {0}, payload size {0}")]
    CannotFitSignature(usize, usize),
    #[error("Invalid signature size {0}")]
    InvalidSignatureSize(usize),
    #[error("Invalid signature")]
    InvalidSignature,
    #[error("Invalid signing/verifying key")]
    InvalidKey,
}

pub trait SignatureKeyStore {
    fn accept_public_key(&mut self, algorithm: u8, data: &[u8]) -> Option<VerificationMethod>;
    fn signer_for_algorithm(&mut self, algorithm: u8) -> Option<SignatureMethod>;

    fn offer_signature_algorithms(&self) -> &[u8];

    fn signer(&mut self, algorithms: &[u8]) -> Option<SignatureMethod> {
        for &algo in algorithms {
            if let Some(signer) = self.signer_for_algorithm(algo) {
                return Some(signer);
            }
        }
        None
    }
}

impl VerifyEd25519 {
    pub fn from_slice(key: &[u8]) -> Result<Self, Error> {
        if key.len() != ed25519_dalek::PUBLIC_KEY_LENGTH {
            return Err(Error::InvalidKey);
        }
        let mut bytes = [0; 32];
        bytes.copy_from_slice(key);
        let verifying_key =
            ed25519_dalek::VerifyingKey::from_bytes(&bytes).map_err(|_| Error::InvalidKey)?;
        Ok(Self { verifying_key })
    }

    pub fn verify(&mut self, message: &[u8], signature: &[u8]) -> Result<(), Error> {
        if signature.len() != ed25519_dalek::SIGNATURE_LENGTH {
            return Err(Error::InvalidSignatureSize(signature.len()));
        }
        let signature = ed25519_dalek::Signature::from_slice(signature).unwrap();
        self.verifying_key
            .verify_strict(message, &signature)
            .map_err(|_| Error::InvalidSignature)
    }

    pub fn fingerprint(&self) -> String {
        fingerprint_sha256("ed25519", self.verifying_key.as_bytes())
    }
}

impl SignEd25519 {
    pub fn generate() -> Result<Self, Error> {
        let mut rng = rand::thread_rng();
        let signing_key = ed25519_dalek::SigningKey::generate(&mut rng);
        let verifying_key = signing_key.verifying_key();
        Ok(Self {
            signing_key,
            verifying_key,
        })
    }

    pub fn load_signing_key<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let signing_key =
            ed25519_dalek::SigningKey::read_pkcs8_pem_file(path).map_err(|_| Error::InvalidKey)?;
        let verifying_key = signing_key.verifying_key();
        Ok(Self {
            signing_key,
            verifying_key,
        })
    }

    pub fn sign(&mut self, message: &[u8], buffer: &mut [u8]) -> Result<usize, Error> {
        if buffer.len() < ed25519_dalek::SIGNATURE_LENGTH {
            return Err(Error::CannotFitSignature(
                buffer.len(),
                ed25519_dalek::SIGNATURE_LENGTH,
            ));
        }
        let signature = self.signing_key.sign(message);
        buffer[..ed25519_dalek::SIGNATURE_LENGTH].copy_from_slice(&signature.to_bytes());
        Ok(ed25519_dalek::SIGNATURE_LENGTH)
    }

    pub fn verifying_key_bytes(&self) -> Cow<[u8]> {
        Cow::Borrowed(self.verifying_key.as_bytes())
    }

    pub fn fingerprint(&self) -> String {
        fingerprint_sha256("ed25519", self.verifying_key.as_bytes())
    }
}

impl VerificationMethod {
    pub fn from_parts(algorithm: u8, data: &[u8]) -> Result<Self, Error> {
        match algorithm {
            V1_SIG_ED25519 => VerifyEd25519::from_slice(data).map(Self::Ed25519),
            _ => todo!(),
        }
    }

    pub fn verify(&mut self, message: &[u8], signature: &[u8]) -> Result<(), Error> {
        match self {
            Self::Ed25519(ed25519) => ed25519.verify(message, signature),
        }
    }

    pub fn fingerprint(&self) -> String {
        match self {
            Self::Ed25519(ed25519) => ed25519.fingerprint(),
        }
    }
}

impl SignatureMethod {
    pub fn generate(algorithm: u8) -> Result<Self, Error> {
        match algorithm {
            V1_SIG_ED25519 => SignEd25519::generate().map(Self::Ed25519),
            _ => todo!(),
        }
    }

    pub fn sign(&mut self, message: &[u8], signature: &mut [u8]) -> Result<usize, Error> {
        match self {
            Self::Ed25519(ed25519) => ed25519.sign(message, signature),
        }
    }

    pub fn verifying_key_bytes(&self) -> Cow<[u8]> {
        match self {
            Self::Ed25519(ed25519) => ed25519.verifying_key_bytes(),
        }
    }

    pub fn algorithm(&self) -> u8 {
        match self {
            Self::Ed25519(_) => V1_SIG_ED25519,
        }
    }

    pub fn fingerprint(&self) -> String {
        match self {
            Self::Ed25519(ed25519) => ed25519.fingerprint(),
        }
    }
}

pub fn fingerprint_sha256(key_algo: &str, key_bytes: &[u8]) -> String {
    let mut digest = sha2::Sha256::new();
    digest.write_all(key_bytes).unwrap();
    let hash = digest.finalize();
    let mut result = format!("{} sha256 ", key_algo);
    for (i, &byte) in hash.iter().enumerate() {
        if i != 0 {
            result.push(':');
        }
        result.push_str(&format!("{byte:02x}"));
    }
    result
}
