use std::{collections::HashSet, path::PathBuf};

use crate::crypt::{
    symmetric::{V1_CIPHER_AES_128_GCM, V1_CIPHER_AES_256_GCM, V1_CIPHER_CHACHA20POLY1305},
    V1_KEX_X25519_DALEK,
};

use super::{
    sig_algo_name,
    signature::{
        fingerprint_sha256, SignEd25519, SignatureKeyStore, SignatureMethod, VerificationMethod,
    },
    V1_SIG_ED25519,
};

fn default_select_kex_algorithm(offer: &[u8]) -> Option<u8> {
    const ACCEPTED: &[u8] = &[V1_KEX_X25519_DALEK];

    for accepted in ACCEPTED {
        if offer.contains(accepted) {
            return Some(*accepted);
        }
    }
    None
}

fn default_select_ciphersuite(offer: &[u8]) -> Option<u8> {
    const ACCEPTED: &[u8] = &[
        V1_CIPHER_CHACHA20POLY1305,
        V1_CIPHER_AES_256_GCM,
        V1_CIPHER_AES_128_GCM,
    ];

    for accepted in ACCEPTED {
        if offer.contains(accepted) {
            return Some(*accepted);
        }
    }
    None
}

fn default_offer_sig_algorithms() -> &'static [u8] {
    &[V1_SIG_ED25519]
}

fn default_offer_kex_algorithms() -> &'static [u8] {
    &[V1_KEX_X25519_DALEK]
}

fn default_offer_ciphersuites() -> &'static [u8] {
    &[
        V1_CIPHER_CHACHA20POLY1305,
        V1_CIPHER_AES_256_GCM,
        V1_CIPHER_AES_128_GCM,
    ]
}

pub struct ClientConfig {
    pub signature_keystore: Box<dyn SignatureKeyStore>,
    pub select_kex_algorithm: fn(&[u8]) -> Option<u8>,
    pub select_ciphersuite: fn(&[u8]) -> Option<u8>,
}

pub struct ServerConfig {
    pub signature_keystore: Box<dyn SignatureKeyStore>,
    pub offer_sig_algorhtms: fn() -> &'static [u8],
    pub offer_kex_algorithms: fn() -> &'static [u8],
    pub offer_ciphersuites: fn() -> &'static [u8],
}

pub struct SimpleServerKeyStore {
    pub path: PathBuf,
    pub accepted_keys: HashSet<String>,
}

pub struct SimpleClientKeyStore {
    algo: [u8; 1],
    key: Option<SignatureMethod>,
}

impl SignatureKeyStore for SimpleServerKeyStore {
    fn offer_signature_algorithms(&self) -> &[u8] {
        &[V1_SIG_ED25519]
    }

    fn signer_for_algorithm(&mut self, algorithm: u8) -> Option<SignatureMethod> {
        match algorithm {
            V1_SIG_ED25519 => SignEd25519::load_signing_key(self.path.join("id_ed25519"))
                .map(SignatureMethod::Ed25519)
                .ok(),
            _ => unreachable!(),
        }
    }

    fn accept_public_key(&mut self, algorithm: u8, data: &[u8]) -> Option<VerificationMethod> {
        let algorithm_name = sig_algo_name(algorithm)?;
        let fingerprint = fingerprint_sha256(algorithm_name, data);
        if self.accepted_keys.contains(&fingerprint) {
            VerificationMethod::from_parts(algorithm, data).ok()
        } else {
            None
        }
    }
}

impl SimpleClientKeyStore {
    pub fn new(key: SignatureMethod) -> Self {
        Self {
            algo: match key {
                SignatureMethod::Ed25519(_) => [V1_SIG_ED25519],
            },
            key: Some(key),
        }
    }
}

impl SignatureKeyStore for SimpleClientKeyStore {
    fn signer_for_algorithm(&mut self, algorithm: u8) -> Option<SignatureMethod> {
        if algorithm == self.algo[0] {
            self.key.take()
        } else {
            None
        }
    }

    fn accept_public_key(&mut self, algorithm: u8, data: &[u8]) -> Option<VerificationMethod> {
        // Accept any key
        VerificationMethod::from_parts(algorithm, data).ok()
    }

    fn offer_signature_algorithms(&self) -> &[u8] {
        &self.algo
    }
}

impl ServerConfig {
    pub fn with_default_algorithms<K: SignatureKeyStore + 'static>(signature_keystore: K) -> Self {
        Self {
            offer_ciphersuites: default_offer_ciphersuites,
            offer_sig_algorhtms: default_offer_sig_algorithms,
            offer_kex_algorithms: default_offer_kex_algorithms,
            signature_keystore: Box::new(signature_keystore),
        }
    }
}

impl ClientConfig {
    pub fn with_default_algorithms<K: SignatureKeyStore + 'static>(signature_keystore: K) -> Self {
        Self {
            select_kex_algorithm: default_select_kex_algorithm,
            select_ciphersuite: default_select_ciphersuite,
            signature_keystore: Box::new(signature_keystore),
        }
    }
}
