// SPDX-FileCopyrightText: 2025 Greenbone AG
//
// SPDX-License-Identifier: GPL-2.0-or-later WITH x11vnc-openssl-exception

use std::{collections::HashMap, sync::RwLock};

use itertools::Itertools;
use kb::InMemoryKbStorage;

use greenbone_scanner_framework::models::FeedType;

use crate::notus::advisories::VulnerabilityData;

use super::{
    Dispatcher, Remover, Retriever, ScanID,
    error::StorageError,
    items::{
        kb::{GetKbContextKey, KbContextKey, KbItem},
        notus_advisory::NotusAdvisory,
        nvt::{Feed, FeedVersion, FileName, Oid},
        result::{ResultContextKeySingle, ResultItem},
    },
};
use greenbone_scanner_framework::models::VTData;

pub mod kb;

/// Contains the FeedType and OID
///
/// The feedtype is required to identify which entries need to be deleted before an feedupate.
///
/// For an example when an Notus feed is updated we need to delete the Notus Vt entries without touching the NASL
/// entries.
type VTKey = (FeedType, String);

/// Vts are using a relative file path as a key. This should make includes, script_dependency
/// lookups relative simple.
type Vts = HashMap<VTKey, VTData>;

/// The results generated by log_, security_, error_message.
type Results = HashMap<ScanID, Vec<ResultItem>>;

// Artificial type for fetching all OIDs
struct OIDs;

/// Is a in-memory dispatcher that behaves like a Storage.
#[derive(Default, Debug)]
pub struct InMemoryStorage {
    vts: RwLock<Vts>,
    oid_lookup: RwLock<HashMap<VTKey, String>>,
    feed_version: RwLock<String>,
    kbs: InMemoryKbStorage,
    results: RwLock<Results>,
}

impl InMemoryStorage {
    /// Creates a new InMemoryStorage.
    pub fn new() -> Self {
        Self {
            ..Default::default()
        }
    }
    fn to_nasl_key(oid: &str) -> (FeedType, String) {
        (FeedType::NASL, oid.to_string())
    }

    /// Stores an already existing Vts structure.
    pub fn set_vts(&self, vts: HashMap<String, VTData>) -> Result<(), StorageError> {
        let mut vts_internal = self.vts.write()?;
        let mut oid_lookup = self.oid_lookup.write()?;
        let mut new_oid_lookup = HashMap::with_capacity(oid_lookup.len());

        for (filename, nvt) in &vts {
            new_oid_lookup.insert(Self::to_nasl_key(&nvt.oid), filename.clone());
        }
        for ((ft, k), v) in oid_lookup.clone().into_iter() {
            if ft == FeedType::Products || ft == FeedType::Advisories {
                new_oid_lookup.insert((ft, k), v);
            }
        }
        *oid_lookup = new_oid_lookup;

        *vts_internal = vts
            .into_iter()
            .map(|(k, v)| (Self::to_nasl_key(&k), v))
            .collect();

        Ok(())
    }

    fn to_notus_advisory_key(oid: &str) -> (FeedType, String) {
        (FeedType::Advisories, oid.to_string())
    }

    fn cache_notus_advisory(&self, adv: VulnerabilityData) -> Result<(), StorageError> {
        let item: VTData = adv.into();

        let mut oid_lookup = self.oid_lookup.write()?;
        oid_lookup.insert(
            Self::to_notus_advisory_key(&item.oid),
            item.filename.clone(),
        );
        drop(oid_lookup); // this drop is necessary because it holds a lock on &self
        let mut vts = self.vts.write()?;
        vts.insert(Self::to_notus_advisory_key(&item.filename), item);

        Ok(())
    }

    fn all_vts(&self) -> Result<Vec<VTData>, StorageError> {
        Ok(self.vts.read()?.values().cloned().collect())
    }

    /// Removes all stored nasl_vts
    fn clean_vts(&self) -> Result<(), StorageError> {
        let mut vts = self.vts.write()?;
        vts.retain(|(ft, _), _| matches!(ft, FeedType::Advisories));
        let mut version = self.feed_version.write()?;
        *version = String::new();
        Ok(())
    }
}

impl Dispatcher<KbContextKey> for InMemoryStorage {
    type Item = KbItem;
    fn dispatch(&self, key: KbContextKey, item: Self::Item) -> Result<(), StorageError> {
        self.kbs.dispatch(key, item)
    }
}

impl Retriever<KbContextKey> for InMemoryStorage {
    type Item = Vec<KbItem>;
    fn retrieve(&self, key: &KbContextKey) -> Result<Option<Self::Item>, StorageError> {
        self.kbs.retrieve(key)
    }
}

impl Retriever<GetKbContextKey> for InMemoryStorage {
    type Item = Vec<(String, Vec<KbItem>)>;
    fn retrieve(&self, key: &GetKbContextKey) -> Result<Option<Self::Item>, StorageError> {
        self.kbs.retrieve(key)
    }
}

impl Remover<KbContextKey> for InMemoryStorage {
    type Item = Vec<KbItem>;
    fn remove(&self, key: &KbContextKey) -> Result<Option<Self::Item>, StorageError> {
        self.kbs.remove(key)
    }
}

impl Dispatcher<()> for InMemoryStorage {
    type Item = NotusAdvisory;
    fn dispatch(&self, _: (), item: Self::Item) -> Result<(), StorageError> {
        self.cache_notus_advisory(item)
    }
}

impl Dispatcher<FileName> for InMemoryStorage {
    type Item = VTData;
    /// Dispatch a single NVT into the storage with a given Key
    fn dispatch(&self, key: FileName, item: Self::Item) -> Result<(), StorageError> {
        let mut vts = self.vts.write()?;
        let mut oid_lookup = self.oid_lookup.write()?;
        oid_lookup.insert(Self::to_nasl_key(&item.oid), key.0.clone());
        vts.insert(Self::to_nasl_key(&key.0), item);
        Ok(())
    }
}

impl Dispatcher<FeedVersion> for InMemoryStorage {
    type Item = String;
    /// Dispatch the feed version into the storage
    fn dispatch(&self, _: FeedVersion, item: Self::Item) -> Result<(), StorageError> {
        let mut feed_version = self.feed_version.write()?;
        *feed_version = item;
        Ok(())
    }
}

impl Retriever<FeedVersion> for InMemoryStorage {
    type Item = String;
    /// Retrieve the feed version from the storage
    fn retrieve(&self, _: &FeedVersion) -> Result<Option<Self::Item>, StorageError> {
        Ok(Some(self.feed_version.read()?.clone()))
    }
}

impl Retriever<Feed> for InMemoryStorage {
    type Item = Vec<VTData>;
    /// Retrieve all NVTs from the storage
    fn retrieve(&self, _: &Feed) -> Result<Option<Self::Item>, StorageError> {
        self.all_vts().map(Some)
    }
}

impl Retriever<FileName> for InMemoryStorage {
    type Item = VTData;
    fn retrieve(&self, key: &FileName) -> Result<Option<Self::Item>, StorageError> {
        let vts = self.vts.read()?;
        // Notus is favored when available. This is done to prevent duplicate definitions.
        Ok(
            if let Some(notus_result) = vts.get(&Self::to_notus_advisory_key(&key.0)) {
                Some(notus_result.clone())
            } else {
                vts.get(&Self::to_nasl_key(&key.0)).cloned()
            },
        )
    }
}

impl Retriever<Oid> for InMemoryStorage {
    type Item = VTData;
    fn retrieve(&self, key: &Oid) -> Result<Option<Self::Item>, StorageError> {
        let oid_lookup = self.oid_lookup.read()?;
        let lookup = |key| match oid_lookup.get(key) {
            None => Ok(None),
            Some(file_name) => {
                let vts = self.vts.read()?;
                let (ft, _) = key;
                Ok(vts.get(&(*ft, file_name.to_string())).cloned())
            }
        };
        match lookup(&Self::to_notus_advisory_key(&key.0)) {
            Ok(Some(vt)) => Ok(Some(vt)),
            Err(e) => Err(e),
            _ => lookup(&Self::to_nasl_key(&key.0)),
        }
    }
}

impl Retriever<OIDs> for InMemoryStorage {
    type Item = Vec<String>;
    /// Retrieve all OIDs from the storage
    fn retrieve(&self, _: &OIDs) -> Result<Option<Self::Item>, StorageError> {
        let vts = self.oid_lookup.read()?;

        let vts = vts.keys().map(|(_, oid)| oid.to_string()).collect_vec();
        Ok(Some(vts))
    }
}

impl Remover<Feed> for InMemoryStorage {
    type Item = ();
    fn remove(&self, _: &Feed) -> Result<Option<Self::Item>, StorageError> {
        self.clean_vts()?;
        Ok(Some(()))
    }
}

impl Dispatcher<ScanID> for InMemoryStorage {
    type Item = ResultItem;
    fn dispatch(&self, key: ScanID, item: Self::Item) -> Result<(), StorageError> {
        let mut results = self.results.write()?;
        match results.get_mut(&key) {
            Some(scan_results) => {
                scan_results.push(item);
            }
            _ => {
                results.insert(key, vec![item]);
            }
        }
        Ok(())
    }
}

impl Retriever<ResultContextKeySingle> for InMemoryStorage {
    type Item = ResultItem;
    fn retrieve(&self, key: &ResultContextKeySingle) -> Result<Option<Self::Item>, StorageError> {
        let results = self.results.read()?;
        if let Some(scan_results) = results.get(&key.0) {
            return Ok(scan_results.get(key.1).cloned());
        }
        Ok(None)
    }
}

impl Retriever<ScanID> for InMemoryStorage {
    type Item = Vec<ResultItem>;
    fn retrieve(&self, key: &ScanID) -> Result<Option<Self::Item>, StorageError> {
        let results = self.results.read()?;

        Ok(results.get(key).cloned())
    }
}

impl Remover<ScanID> for InMemoryStorage {
    type Item = Vec<ResultItem>;
    fn remove(&self, key: &ScanID) -> Result<Option<Self::Item>, StorageError> {
        let mut results = self.results.write()?;
        Ok(results.remove(key))
    }
}

impl Remover<ResultContextKeySingle> for InMemoryStorage {
    type Item = ResultItem;
    fn remove(&self, key: &ResultContextKeySingle) -> Result<Option<ResultItem>, StorageError> {
        let mut results = self.results.write()?;
        if let Some(results) = results.get_mut(&key.0) {
            return Ok(Some(results.remove(key.1)));
        }
        Ok(None)
    }
}

#[cfg(test)]
mod tests {
    use crate::storage::{
        Dispatcher, Retriever,
        error::StorageError,
        inmemory::InMemoryStorage,
        items::{
            kb::{KbContextKey, KbItem},
            nvt::{FileName, Oid},
        },
    };
    use greenbone_scanner_framework::models::VTData;

    #[test]
    fn nvt() -> Result<(), StorageError> {
        let storage = InMemoryStorage::default();
        let key = FileName(String::new());
        let nvt = VTData {
            oid: "moep".to_string(),
            ..VTData::default()
        };
        storage.dispatch(key.clone(), nvt.clone())?;
        let ret = storage.retrieve(&key).unwrap().unwrap();
        assert_eq!(nvt, ret);
        Ok(())
    }

    #[test]
    fn nvt_oid() -> Result<(), StorageError> {
        let storage = InMemoryStorage::default();
        let key = FileName(String::new());
        let nvt = VTData {
            oid: "moep".to_string(),
            ..VTData::default()
        };
        storage.dispatch(key.clone(), nvt.clone())?;
        let key = Oid(nvt.oid.clone());
        let ret = storage.retrieve(&key).unwrap().unwrap();
        assert_eq!(nvt, ret);
        Ok(())
    }

    #[test]
    fn kb() {
        let storage = InMemoryStorage::default();
        let key = KbContextKey::default();
        let value1 = KbItem::String("1".to_string());
        let value2 = KbItem::String("2".to_string());
        storage.dispatch(key.clone(), value1.clone()).unwrap();
        storage.dispatch(key.clone(), value2.clone()).unwrap();
        let ret = storage.retrieve(&key).unwrap().unwrap();
        assert_eq!(ret, vec![value1, value2]);
    }
}
