use std::{
    io::{self, Error, ErrorKind},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use tokio::sync::RwLock;

pub struct Database {
    inner: Arc<RwLock<Box<dyn crate::rpcchainvm::database::Database + Send + Sync>>>,
    corrupted: AtomicBool,
    corrupted_error: Arc<RwLock<String>>,
}

impl Database {
    pub fn new(
        db: Box<dyn crate::rpcchainvm::database::Database + Send + Sync>,
    ) -> Box<dyn crate::rpcchainvm::database::Database + Send + Sync> {
        Box::new(Database {
            inner: Arc::new(RwLock::new(db)),
            corrupted: AtomicBool::new(false),
            corrupted_error: Arc::new(RwLock::new("".into())),
        })
    }
}

/// Database is a wrapper around Database it prevents any future calls
/// in case a corruption occurs.
/// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/database/corruptabledb#Database
impl crate::rpcchainvm::database::Database for Database {}

#[tonic::async_trait]
impl crate::rpcchainvm::database::KeyValueReaderWriterDeleter for Database {
    /// Attempts to return if the database has a key with the provided value.
    async fn has(&self, key: &[u8]) -> io::Result<bool> {
        if self.corrupted.load(Ordering::Relaxed) {
            return Err(error_from_string(
                self.corrupted_error.read().await.to_string(),
            ));
        }

        let db = self.inner.read().await;
        match db.get(key).await {
            Ok(_) => Ok(true),
            Err(err) => {
                let (is_corrupted, err) = handle_error(err).await;
                if is_corrupted {
                    *self.corrupted_error.write().await = err.to_string();
                    self.corrupted.store(true, Ordering::Relaxed);
                    return Err(error_from_string(
                        self.corrupted_error.read().await.to_string(),
                    ));
                }
                if err.kind() == ErrorKind::NotFound {
                    return Ok(false);
                }
                return Err(err);
            }
        }
    }

    /// Attempts to return the value that was mapped to the key that was provided.
    async fn get(&self, key: &[u8]) -> io::Result<Vec<u8>> {
        if self.corrupted.load(Ordering::Relaxed) {
            return Err(error_from_string(
                self.corrupted_error.read().await.to_string(),
            ));
        }

        let db = self.inner.read().await;
        match db.get(key).await {
            Ok(resp) => Ok(resp),
            Err(err) => {
                let (is_corrupted, err) = handle_error(err).await;
                if is_corrupted {
                    *self.corrupted_error.write().await = err.to_string();
                    self.corrupted.store(true, Ordering::Relaxed);
                    return Err(error_from_string(
                        self.corrupted_error.read().await.to_string(),
                    ));
                }
                return Err(err);
            }
        }
    }

    /// Attempts to set the value this key maps to.
    async fn put(&mut self, key: &[u8], value: &[u8]) -> io::Result<()> {
        if self.corrupted.load(Ordering::Relaxed) {
            return Err(error_from_string(
                self.corrupted_error.read().await.to_string(),
            ));
        }

        let mut db = self.inner.write().await;
        match db.put(key, value).await {
            Ok(_) => Ok(()),
            Err(err) => {
                let (is_corrupted, err) = handle_error(err).await;
                if is_corrupted {
                    *self.corrupted_error.write().await = err.to_string();
                    self.corrupted.store(true, Ordering::Relaxed);
                    return Err(error_from_string(
                        self.corrupted_error.read().await.to_string(),
                    ));
                }
                return Err(err);
            }
        }
    }

    /// Attempts to remove any mapping from the key.
    async fn delete(&mut self, key: &[u8]) -> io::Result<()> {
        if self.corrupted.load(Ordering::Relaxed) {
            return Err(error_from_string(
                self.corrupted_error.read().await.to_string(),
            ));
        }

        let mut db = self.inner.write().await;
        match db.delete(key).await {
            Ok(_) => Ok(()),
            Err(err) => {
                let (is_corrupted, err) = handle_error(err).await;
                if is_corrupted {
                    *self.corrupted_error.write().await = err.to_string();
                    self.corrupted.store(true, Ordering::Relaxed);
                    return Err(error_from_string(
                        self.corrupted_error.read().await.to_string(),
                    ));
                }
                return Err(err);
            }
        }
    }
}

#[tonic::async_trait]
impl crate::rpcchainvm::database::Closer for Database {
    /// Attempts to close the database.
    async fn close(&self) -> io::Result<()> {
        if self.corrupted.load(Ordering::Relaxed) {
            return Err(error_from_string(
                self.corrupted_error.read().await.to_string(),
            ));
        }

        let db = self.inner.write().await;
        match db.close().await {
            Ok(_) => Ok(()),
            Err(err) => {
                let (is_corrupted, err) = handle_error(err).await;
                if is_corrupted {
                    *self.corrupted_error.write().await = err.to_string();
                    self.corrupted.store(true, Ordering::Relaxed);
                    return Err(error_from_string(
                        self.corrupted_error.read().await.to_string(),
                    ));
                }
                return Err(err);
            }
        }
    }
}

#[tonic::async_trait]
impl crate::rpcchainvm::health::Checkable for Database {
    /// Checks if the database has been closed.
    async fn health_check(&self) -> io::Result<Vec<u8>> {
        if self.corrupted.load(Ordering::Relaxed) {
            return Err(error_from_string(
                self.corrupted_error.read().await.to_string(),
            ));
        }

        let db = self.inner.read().await;
        match db.health_check().await {
            Ok(resp) => Ok(resp),
            Err(err) => {
                let (is_corrupted, err) = handle_error(err).await;
                if is_corrupted {
                    *self.corrupted_error.write().await = err.to_string();
                    self.corrupted.store(true, Ordering::Relaxed);
                    return Err(error_from_string(
                        self.corrupted_error.read().await.to_string(),
                    ));
                }
                return Err(err);
            }
        }
    }
}

/// Accepts an error and returns a corruption error if the original error is not "database closed"
/// or "not found".
async fn handle_error(error: Error) -> (bool, Error) {
    match error {
        e if e.to_string().contains("database closed") => return (false, e),
        e if e.kind() == ErrorKind::NotFound && e.to_string().contains("not found") => {
            return (false, e)
        }
        _ => (
            true,
            Error::new(
                ErrorKind::Other,
                format!(
                    "closed to avoid possible corruption, init error: {}",
                    error.to_string()
                ),
            ),
        ),
    }
}

fn error_from_string(message: String) -> Error {
    return Error::new(ErrorKind::Other, message);
}
