// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::{Arc, Mutex};

use tracing::instrument;
use x509_cert::der::Encode;
use yubikey::piv::AlgorithmId;
use yubikey::YubiKey;

use crate::piv::slot::SlotObject;
use crate::piv::util::get_algorithm_id_by_certificate;
use native_pkcs11_traits::Certificate as P11Certificate;
use native_pkcs11_traits::KeyAlgorithm as P11KeyAlgorithm;
use native_pkcs11_traits::KeySearchOptions as P11KeySearchOptions;
use native_pkcs11_traits::PrivateKey as P11PrivateKey;
use native_pkcs11_traits::PublicKey as P11PublicKey;
use native_pkcs11_traits::Result as P11Result;
use native_pkcs11_traits::{Backend, KeyAlgorithm, KeySearchOptions};

#[derive(Debug, Default)]
pub struct YubikeyPivBackend {
    cached_pin: Mutex<Option<String>>,
    yubikey: Mutex<Option<YubiKey>>,
    slot_objects: Mutex<Option<Vec<SlotObject>>>,
}

impl YubikeyPivBackend {
    pub fn new() -> Self {
        YubikeyPivBackend::default()
    }

    fn run_with_yubikey<F>(&self, verify: bool, mut callback: F) -> P11Result<()>
    where
        F: FnMut(&mut YubiKey) -> P11Result<()>,
    {
        let mut yubikey = self.yubikey.lock().unwrap();
        if yubikey.is_none() {
            *yubikey = Some(YubiKey::open()?);
        }
        let mut yk = yubikey.as_mut().unwrap();
        if verify {
            let pin = self.prepare_pin()?;
            let verify_result = yk.verify_pin(pin.as_bytes());
            if verify_result.is_err() {
                self.clear_pin();
            }
            verify_result?;
        }
        callback(&mut yk)
    }

    fn clear_pin(&self) -> () {
        let mut cached_pin = self.cached_pin.lock().unwrap();
        if cached_pin.is_some() {
            *cached_pin = None;
        }
    }

    fn prepare_pin(&self) -> P11Result<String> {
        let mut cached_pin = self.cached_pin.lock().unwrap();
        if cached_pin.is_none() {
            let pin = crate::piv::pinentry::get_pin()?;
            *cached_pin = Some(pin);
        }
        Ok(cached_pin.as_deref().unwrap().to_string())
    }

    fn init_slot_objects(&self) -> P11Result<()> {
        let mut slot_objects_opt = self.slot_objects.lock().unwrap();
        if slot_objects_opt.is_some() {
            return Ok(());
        }

        let mut slot_objects = vec![];
        self.run_with_yubikey(false, |yk| {
            let keys = yk.piv_keys()?;
            for key in keys {
                let slot_id = key.slot().to_string();
                let certificate_der = key.certificate().cert.to_der()?;
                let public_key_der = key.certificate().cert.tbs_certificate.subject_public_key_info.to_der()?;

                let algorithm_id = get_algorithm_id_by_certificate(key.certificate())?;
                let algorithm = match algorithm_id {
                    AlgorithmId::Rsa1024 | AlgorithmId::Rsa2048 => KeyAlgorithm::Rsa,
                    AlgorithmId::EccP256 | AlgorithmId::EccP384 => KeyAlgorithm::Ecc,
                };
                // println!(">>> {} {:?} {}", &slot_id, algorithm_id, &key.certificate().cert.tbs_certificate.subject); // TODO remove
                let slot_object = SlotObject::new(
                    algorithm,
                    slot_id,
                    certificate_der,
                    public_key_der,
                );
                slot_objects.push(slot_object);
            }
            Ok(())
        })?;
        *slot_objects_opt = Some(slot_objects);
        Ok(())
    }

    fn for_each_slot_objects<F>(&self, mut callback: F) -> P11Result<()>
    where
        F: FnMut(&SlotObject) -> P11Result<()>,
    {
        let slot_objects = &self.slot_objects.lock().unwrap();
        for slot_objects in slot_objects.iter() {
            for slot_object in slot_objects {
                callback(slot_object)?;
            }
        }
        Ok(())
    }
}

impl Backend for YubikeyPivBackend {
    fn name(&self) -> String {
        "Yubikey PIV".into()
    }

    #[instrument]
    fn find_all_certificates(
        &self,
    ) -> P11Result<Vec<Box<dyn P11Certificate>>> {
        // println!("[find_all_certificates]");
        self.init_slot_objects()?;
        let mut certs: Vec<Box<dyn P11Certificate>> = vec![];
        self.for_each_slot_objects(|slot_object| {
            certs.push(slot_object.to_certificate()?);
            Ok(())
        })?;
        Ok(certs)
    }

    #[instrument]
    fn find_private_key(
        &self,
        query: P11KeySearchOptions,
    ) -> P11Result<Option<Arc<dyn P11PrivateKey>>> {
        println!("[find_private_key]");
        self.init_slot_objects()?;
        let mut private_key: Option<Arc<dyn P11PrivateKey>> = None;
        self.for_each_slot_objects(|slot_object| {
            if private_key.is_none() {
                match &query {
                    KeySearchOptions::Label(label) => {
                        if &slot_object.label == label {
                            private_key = Some(Arc::from(slot_object.to_private_key()?));
                        }
                    }
                    KeySearchOptions::PublicKeyHash(public_key_hash) => {
                        if hex::encode(&slot_object.public_key_hash) == hex::encode(public_key_hash) {
                            private_key = Some(Arc::from(slot_object.to_private_key()?));
                        }
                    }
                }
            }
            Ok(())
        })?;
        Ok(private_key)
    }

    #[instrument]
    fn find_public_key(
        &self,
        query: P11KeySearchOptions,
    ) -> P11Result<Option<Box<dyn P11PublicKey>>> {
        println!("[find_public_key]");
        self.init_slot_objects()?;
        let mut public_key: Option<Box<dyn P11PublicKey>> = None;
        self.for_each_slot_objects(|slot_object| {
            if public_key.is_none() {
                match &query {
                    KeySearchOptions::Label(label) => {
                        if &slot_object.label == label {
                            public_key = Some(slot_object.to_public_key()?);
                        }
                    }
                    KeySearchOptions::PublicKeyHash(public_key_hash) => {
                        if hex::encode(&slot_object.public_key_hash) == hex::encode(public_key_hash) {
                            public_key = Some(slot_object.to_public_key()?);
                        }
                    }
                }
            }
            Ok(())
        })?;
        Ok(public_key)
    }

    fn find_all_private_keys(
        &self,
    ) -> P11Result<Vec<Arc<dyn P11PrivateKey>>> {
        println!("[find_all_private_keys]");
        self.init_slot_objects()?;
        let mut private_keys: Vec<Arc<dyn P11PrivateKey>> = vec![];
        self.for_each_slot_objects(|slot_object| {
            private_keys.push(Arc::from(slot_object.to_private_key()?));
            Ok(())
        })?;
        Ok(private_keys)
    }

    fn find_all_public_keys(
        &self,
    ) -> P11Result<Vec<Arc<dyn P11PublicKey>>> {
        println!("[find_all_public_keys]");
        self.init_slot_objects()?;
        let mut public_keys: Vec<Arc<dyn P11PublicKey>> = vec![];
        self.for_each_slot_objects(|slot_object| {
            public_keys.push(Arc::from(slot_object.to_public_key()?));
            Ok(())
        })?;
        Ok(public_keys)
    }

    #[instrument]
    fn generate_key(
        &self,
        _algorithm: P11KeyAlgorithm,
        _label: Option<&str>,
    ) -> P11Result<Arc<dyn P11PrivateKey>> {
        Err("Generate key not supported, please use ykman, URL: https://hatter.in/ykman")?
    }
}