use malwaredb_lzjd::{LZDict, Murmur3HashState};
use malwaredb_types::exec::pe32::EXE;

use std::fmt::{Debug, Formatter};
use std::path::Path;

use anyhow::{bail, Context, Result};
use base64::engine::general_purpose;
use base64::Engine;
use fuzzyhash::FuzzyHash;
use home::home_dir;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tlsh_fixed::TlshBuilder;
use tracing::{error, warn};
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Config file name expected by MalwareDB Client
const DOT_MDB_CLIENT_TOML: &str = ".mdb_client.toml";

/// MDB version
pub const MDB_VERSION: &str = env!("CARGO_PKG_VERSION");

/// MalwareDB Client Configuration and connection
#[derive(Deserialize, Serialize, Zeroize, ZeroizeOnDrop)]
pub struct MdbClient {
    /// URL of MalwareDB, including http and port number, ending without a slash
    pub url: String,

    /// User's API key for MalwareDB
    api_key: String,
}

impl MdbClient {
    /// MDB Client from components, doesn't test connectivity
    pub fn new(url: String, api_key: String) -> Self {
        let mut url = url;
        let url = if url.ends_with('/') {
            url.pop();
            url
        } else {
            url
        };

        Self { url, api_key }
    }

    /// Login to a server, optionally save the config file, and return a client object
    pub async fn login(
        url: String,
        username: String,
        password: String,
        save: bool,
    ) -> Result<Self> {
        let mut url = url;
        let url = if url.ends_with('/') {
            url.pop();
            url
        } else {
            url
        };

        let api_request = malwaredb_api::GetAPIKeyRequest {
            user: username,
            password,
        };

        let client = reqwest::Client::new();
        let res = client
            .post(format!("{url}{}", malwaredb_api::USER_LOGIN_URL))
            .json(&api_request)
            .send()
            .await?
            .json::<malwaredb_api::GetAPIKeyResponse>()
            .await?;

        if let Some(key) = &res.key {
            let client = MdbClient {
                url,
                api_key: key.clone(),
            };

            if save {
                if let Err(e) = client.save() {
                    error!("Login successful but failed to save config: {e}");
                    bail!("Login successful but failed to save config: {e}");
                }
            }
            Ok(client)
        } else {
            if let Some(msg) = &res.message {
                error!("Login failed, response: {msg}");
            }
            bail!("server error or bad credentials");
        }
    }

    /// Reset one's own API key to effectively logout & disable all clients who are using the key
    pub async fn reset_key(&self) -> Result<()> {
        let client = reqwest::Client::new();
        let response = client
            .get(format!("{}{}", self.url, malwaredb_api::USER_LOGOUT_URL))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await
            .context("server error, or invalid API key")?;
        if response.status().is_success() {
            bail!("failed to reset API key, was it correct?");
        }
        Ok(())
    }

    /// MDB Client loaded from a specified path
    pub fn from_file(path: &std::path::PathBuf) -> Result<Self> {
        let config = std::fs::read_to_string(path)
            .context(format!("failed to read config file {}", path.display()))?;
        let cfg: MdbClient = toml::from_str(&config)
            .context(format!("failed to parse config file {}", path.display()))?;
        Ok(cfg)
    }

    /// MDB Client from user's home directory
    pub fn load() -> Result<Self> {
        let config = Path::new("mdb_client.toml");
        if config.exists() {
            return Self::from_file(&config.to_path_buf());
        }

        if let Some(mut home_config) = home_dir() {
            home_config.push(DOT_MDB_CLIENT_TOML);
            if home_config.exists() {
                return Self::from_file(&home_config);
            }
        }
        bail!("config file not found")
    }

    /// Save MDB Client to the user's home directory
    pub fn save(&self) -> Result<()> {
        let toml = toml::to_string(self)?;
        if let Some(mut home_config) = home_dir() {
            home_config.push(DOT_MDB_CLIENT_TOML);
            std::fs::write(&home_config, toml).context(format!(
                "Unable to write config file at {}",
                &home_config.display()
            ))?;
            return Ok(());
        }

        std::fs::write("mdb_client.toml", toml).context("failed to write mdb config")
    }

    /// Delete the MalwareDB client config file
    pub fn delete(&self) -> Result<()> {
        if let Some(mut home_config) = home_dir() {
            home_config.push(DOT_MDB_CLIENT_TOML);
            if home_config.exists() {
                std::fs::remove_file(home_config)?;
            }
        }
        Ok(())
    }

    // Actions of the client

    /// Get information about the server, unauthenticated
    pub async fn server_info(&self) -> Result<malwaredb_api::ServerInfo> {
        reqwest::get(format!("{}{}", self.url, malwaredb_api::SERVER_INFO))
            .await?
            .json::<malwaredb_api::ServerInfo>()
            .await
            .context("failed to receive or decode server info")
    }

    /// Get file types supported by the server, unauthenticated
    pub async fn supported_types(&self) -> Result<malwaredb_api::SupportedFileTypes> {
        reqwest::get(format!(
            "{}{}",
            self.url,
            malwaredb_api::SUPPORTED_FILE_TYPES
        ))
        .await?
        .json::<malwaredb_api::SupportedFileTypes>()
        .await
        .context("failed to receive or decode server-supported file types")
    }

    /// Get information about the user
    pub async fn whoami(&self) -> Result<malwaredb_api::GetUserInfoResponse> {
        let client = reqwest::Client::new();
        client
            .get(format!("{}{}", self.url, malwaredb_api::USER_INFO_URL))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::GetUserInfoResponse>()
            .await
            .context("failed to receive or decode user info, or invalid API key")
    }

    /// Get the sample labels known to the server
    pub async fn labels(&self) -> Result<malwaredb_api::Labels> {
        let client = reqwest::Client::new();
        client
            .get(format!("{}{}", self.url, malwaredb_api::LIST_LABELS))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::Labels>()
            .await
            .context("failed to receive or decode available labels, or invalid API key")
    }

    /// Get the sources available to the current user
    pub async fn sources(&self) -> Result<malwaredb_api::Sources> {
        let client = reqwest::Client::new();
        client
            .get(format!("{}{}", self.url, malwaredb_api::LIST_SOURCES))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::Sources>()
            .await
            .context("failed to receive or decode available labels, or invalid API key")
    }

    /// Submit one file to MalwareDB: provide the contents, file name, and source ID
    pub async fn submit(
        &self,
        contents: impl AsRef<[u8]>,
        file_name: &str,
        source_id: u32,
    ) -> Result<bool> {
        let mut hasher = Sha256::new();
        hasher.update(&contents);
        let result = hasher.finalize();

        let encoded = general_purpose::STANDARD.encode(contents);

        let payload = malwaredb_api::NewSample {
            file_name: file_name.to_string(),
            source_id,
            file_contents_b64: encoded,
            sha256: hex::encode(result),
        };

        let client = reqwest::Client::new();
        match client
            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .json(&payload)
            .send()
            .await
        {
            Ok(res) => {
                if !res.status().is_success() {
                    warn!("Code {} sending {file_name}", res.status());
                }
                Ok(res.status().is_success())
            }
            Err(e) => {
                error!("Error sending {file_name}: {e}");
                bail!(e.to_string())
            }
        }
    }

    /// Retrieve sample by hash, optionally in the CaRT format
    pub async fn retrieve(&self, hash: &str, cart: bool) -> Result<Vec<u8>> {
        let api_endpoint = if cart {
            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_CART)
        } else {
            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE)
        };

        let client = reqwest::Client::new();
        let res = client
            .get(format!("{}{api_endpoint}", self.url))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?;

        if !res.status().is_success() {
            bail!("Received code {}", res.status());
        }

        let body = res.bytes().await?;
        Ok(body.to_vec())
    }

    /// Fetch a report for a sample
    pub async fn report(&self, hash: &str) -> Result<malwaredb_api::Report> {
        let client = reqwest::Client::new();
        client
            .get(format!(
                "{}{}/{hash}",
                self.url,
                malwaredb_api::SAMPLE_REPORT
            ))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::Report>()
            .await
            .context("failed to receive or decode sample report, or invalid API key")
    }

    /// Find similar samples in MalwareDB based on the contents of a given file.
    /// This does not submit the sample to MalwareDB.
    pub async fn similar(&self, contents: &[u8]) -> Result<malwaredb_api::SimilarSamplesResponse> {
        let mut hashes = vec![];
        let ssdeep_hash = FuzzyHash::new(contents);
        hashes.push((
            malwaredb_api::SimilarityHashType::SDHash,
            ssdeep_hash.to_string(),
        ));

        let build_hasher = Murmur3HashState::default();
        let lzjd_str =
            LZDict::from_bytes_stream(contents.iter().copied(), &build_hasher).to_string();
        hashes.push((malwaredb_api::SimilarityHashType::LZJD, lzjd_str));

        let mut builder = TlshBuilder::new(
            tlsh_fixed::BucketKind::Bucket256,
            tlsh_fixed::ChecksumKind::ThreeByte,
            tlsh_fixed::Version::Version4,
        );

        builder.update(contents);
        if let Ok(hasher) = builder.build() {
            hashes.push((malwaredb_api::SimilarityHashType::TLSH, hasher.hash()));
        }

        if let Ok(exe) = EXE::from(contents) {
            if let Some(imports) = exe.imports {
                hashes.push((
                    malwaredb_api::SimilarityHashType::ImportHash,
                    hex::encode(imports.hash()),
                ));
                hashes.push((
                    malwaredb_api::SimilarityHashType::FuzzyImportHash,
                    imports.fuzzy_hash(),
                ));
            }
        }

        let request = malwaredb_api::SimilarSamplesRequest { hashes };

        let client = reqwest::Client::new();
        client
            .post(format!("{}{}", self.url, malwaredb_api::SIMILAR_SAMPLES))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .json(&request)
            .send()
            .await?
            .json::<malwaredb_api::SimilarSamplesResponse>()
            .await
            .context("failed to receive or decode similarity response, or invalid API key")
    }
}

impl Debug for MdbClient {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "MDB Client v{MDB_VERSION}: {}", self.url)
    }
}