// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module comprises AWS DynamoDb bucket READ ONLY access to download and parse
//! Audit Proofs. THIS IS JUST A SAMPLE AND IS NOT IMPLEMENTED. This will be rounded
//! out later if support is to be added for dynamo db

use super::ProofIndexCacheOption;
use anyhow::Result;
use async_trait::async_trait;
use aws_config::meta::region::RegionProviderChain;
use aws_sdk_dynamodb as dynamo;
use aws_sdk_dynamodb::model::AttributeValue;
use aws_sdk_dynamodb::Region;
use aws_smithy_http::endpoint::Endpoint;
use clap::Args;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[cfg(test)]
mod test;

use super::EpochSummary;

const MIN_TABLE_CHARS: usize = 3;
const MAX_TABLE_CHARS: usize = 255;
const ALLOWED_TABLE_CHARS: [char; 65] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
    't', 'u', 'v', 'w', 'x', 'y', 'z', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L',
    'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'z', '.', '-', '1', '2', '3',
    '4', '5', '6', '7', '8', '9', '0', '_',
];

fn validate_table_name(s: &str) -> Result<String, String> {
    let str = s.to_string();
    if str.len() < MIN_TABLE_CHARS || str.len() > MAX_TABLE_CHARS {
        return Err(format!(
            "Table name must be between [{}, {}] characters in length. Gave {}",
            MIN_TABLE_CHARS,
            MAX_TABLE_CHARS,
            str.len()
        ));
    }

    for c in str.chars() {
        if !ALLOWED_TABLE_CHARS.iter().any(|v| c == *v) {
            return Err(format!("Character '{c}' is not allowed in table name. Table names must contain lower-case letters, numbers, '-', '_', and '.' only."));
        }
    }
    Ok(str)
}

/// Connect to DynamoDb which holds the "listing" information
/// on audit proofs, but the true blobs are stored in S3. (It's just
/// indexed lookups of S3)
#[derive(Args, Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DynamoDbClapSettings {
    /// The dynamoDb table name
    #[clap(long, short,
        value_parser = validate_table_name)]
    pub table: String,

    /// The AWS region where the dynamo db cluster lives
    #[clap(long, short)]
    pub region: String,

    /// The S3 bucket where the audit proofs are stored
    #[clap(
        long,
        value_parser = super::s3::validate_bucket_name
    )]
    pub bucket: String,

    /// (OPTIONAL) AWS DynamoDb custom endpoint
    #[clap(long, value_parser = super::s3::validate_uri)]
    pub dynamo_endpoint: Option<String>,

    /// (OPTIONAL) AWS S3 bucket custom endpoint
    #[clap(long, value_parser = super::s3::validate_uri)]
    pub s3_endpoint: Option<String>,

    /// (OPTIONAL) AWS Access key for the session
    #[clap(long)]
    pub access_key: Option<String>,

    /// (OPTIONAL) AWS secret key for the session
    #[clap(long)]
    pub secret_key: Option<String>,
}

impl super::CommonStorageClapSettings for DynamoDbClapSettings {
    fn bucket(&self) -> String {
        self.bucket.clone()
    }

    fn region(&self) -> String {
        self.region.clone()
    }

    fn s3_endpoint(&self) -> Option<String> {
        self.s3_endpoint.clone()
    }

    fn access_key(&self) -> Option<String> {
        self.access_key.clone()
    }

    fn secret_key(&self) -> Option<String> {
        self.secret_key.clone()
    }
    fn set_bucket(&mut self, v: String) {
        self.bucket = v;
    }
    fn set_region(&mut self, v: String) {
        self.region = v;
    }
    fn set_s3_endpoint(&mut self, v: Option<String>) {
        self.s3_endpoint = v;
    }
    fn set_access_key(&mut self, v: Option<String>) {
        self.access_key = v;
    }
    fn set_secret_key(&mut self, v: Option<String>) {
        self.secret_key = v;
    }
}

#[derive(Debug)]
pub struct DynamoDbAuditStorage {
    /// The table where the audit proofs are stored
    table: String,
    /// The AWS region
    region: String,
    /// The S3 bucket which contains the underlying proof material
    bucket: String,
    /// (OPTIONAL) AWS DynamoDb custom endpoint
    dynamo_endpoint: Option<String>,
    /// (OPTIONAL) AWS S3 bucket custom endpoint
    s3_endpoint: Option<String>,
    /// (OPTIONAL) AWS Access key for the session
    access_key: Option<String>,
    /// (OPTIONAL) AWS secret key for the session
    secret_key: Option<String>,
    /// The configuration, cached for subsequent accesses
    config: Arc<RwLock<Option<aws_config::SdkConfig>>>,
    /// Cache of epoch summaries, populated with each call to list_blob_keys
    cache: Arc<RwLock<Option<Vec<super::EpochSummary>>>>,
}

impl From<&DynamoDbClapSettings> for DynamoDbAuditStorage {
    fn from(clap: &DynamoDbClapSettings) -> Self {
        Self {
            table: clap.table.clone(),
            region: clap.region.clone(),
            bucket: clap.bucket.clone(),
            dynamo_endpoint: clap.dynamo_endpoint.as_ref().cloned(),
            s3_endpoint: clap.s3_endpoint.as_ref().cloned(),
            access_key: clap.access_key.as_ref().cloned(),
            secret_key: clap.secret_key.as_ref().cloned(),
            config: Arc::new(RwLock::new(None)),
            cache: Arc::new(RwLock::new(None)),
        }
    }
}

impl From<&DynamoDbAuditStorage> for super::s3::S3AuditStorage {
    fn from(clap: &DynamoDbAuditStorage) -> Self {
        Self {
            bucket: clap.bucket.clone(),
            region: clap.region.clone(),
            access_key: clap.access_key.as_ref().cloned(),
            secret_key: clap.secret_key.as_ref().cloned(),
            endpoint: clap.s3_endpoint.as_ref().cloned(),
            config: Arc::new(RwLock::new(None)),
            cache: Arc::new(RwLock::new(None)),
        }
    }
}

impl DynamoDbAuditStorage {
    async fn get_shared_config(&self) -> aws_types::SdkConfig {
        let mut lock = self.config.write().await;
        if let Some(config) = &*lock {
            config.clone()
        } else {
            // Get the shared AWS config
            let region_provider = RegionProviderChain::first_try(Region::new(self.region.clone()))
                .or_default_provider()
                .or_else(Region::new(super::s3::DEFAULT_AWS_REGION));

            let mut shared_config_loader = aws_config::from_env().region(region_provider);

            if let (Some(access_key), Some(secret_key)) = (&self.access_key, &self.secret_key) {
                let credentials = aws_types::Credentials::from_keys(access_key, secret_key, None);
                shared_config_loader = shared_config_loader.credentials_provider(credentials);
            }

            if let Some(endpoint) = &self.dynamo_endpoint {
                // THE URI IS ALREADY VALIDATED BY CLAP ARGS, hence the expect here
                let endpoint_resolver =
                    Endpoint::immutable(endpoint.parse().expect("invalid S3 storage URI"));
                shared_config_loader = shared_config_loader.endpoint_resolver(endpoint_resolver);
            }

            let shared_config = shared_config_loader.load().await;

            *lock = Some(shared_config.clone());
            shared_config
        }
    }

    fn convert_digest(data: &[u8]) -> Result<akd::Digest> {
        akd::hash::try_parse_digest(data).map_err(|msg| anyhow::anyhow!(msg))
    }

    fn parse_summary(
        dynamo_attribute_values: &HashMap<String, AttributeValue>,
    ) -> Result<Option<EpochSummary>> {
        if let (
            Some(AttributeValue::N(epoch)),
            Some(AttributeValue::B(phash)),
            Some(AttributeValue::B(chash)),
            Some(AttributeValue::S(blob_key)),
        ) = (
            dynamo_attribute_values.get("epoch"),
            dynamo_attribute_values.get("previous_hash"),
            dynamo_attribute_values.get("current_hash"),
            dynamo_attribute_values.get("blob"),
        ) {
            let phash_digest = Self::convert_digest(phash.as_ref())?;
            let chash_digest = Self::convert_digest(chash.as_ref())?;
            let audit_blob_name = akd::local_auditing::AuditBlobName {
                epoch: epoch.parse()?,
                previous_hash: phash_digest,
                current_hash: chash_digest,
            };
            Ok(Some(EpochSummary {
                key: blob_key.clone(),
                name: audit_blob_name,
            }))
        } else {
            Ok(None)
        }
    }

    async fn query_proofs(&self) -> dynamo::paginator::ScanPaginator {
        let config = self.get_shared_config().await;
        let client = dynamo::Client::new(&config);
        client
            .scan()
            .table_name(self.table.clone())
            .into_paginator()
    }

    fn process_streamed_response(
        part: Result<dynamo::output::ScanOutput, dynamo::types::SdkError<dynamo::error::ScanError>>,
    ) -> Result<Vec<super::EpochSummary>> {
        let mut partials = vec![];
        match part {
            Err(sdk_err) => {
                return Err(anyhow::anyhow!(
                    "Error executing Query in DynamoDb: {}",
                    sdk_err
                ));
            }
            Ok(query_output) => {
                if let Some(contents) = query_output.items {
                    for map_ref in contents {
                        if let Ok(Some(partial)) = Self::parse_summary(&map_ref) {
                            partials.push(partial);
                        }
                    }
                }
            }
        }
        Ok(partials)
    }
}

#[async_trait]
impl super::AuditProofStorage for DynamoDbAuditStorage {
    fn default_cache_control(&self) -> ProofIndexCacheOption {
        ProofIndexCacheOption::NoCache
    }

    async fn list_proofs(
        &self,
        cache_control: ProofIndexCacheOption,
    ) -> Result<Vec<super::EpochSummary>> {
        use tokio_stream::StreamExt;

        {
            if let ProofIndexCacheOption::UseCache = &cache_control {
                if let Some(cache) = &*self.cache.read().await {
                    return Ok(cache.clone());
                }
            }
        }

        let mut results = vec![];
        let client = self.query_proofs().await;
        let mut stream = client.page_size(1000).send();
        while let Some(result) = stream.next().await {
            let mut maybe_item = Self::process_streamed_response(result)?;
            results.append(&mut maybe_item);
        }

        if !results.is_empty() {
            let mut update = self.cache.write().await;
            *update = Some(results.clone());
        }

        Ok(results)
    }

    async fn get_proof(
        &self,
        epoch: &super::EpochSummary,
    ) -> Result<akd::local_auditing::AuditBlob> {
        let s3_storage: crate::storage::s3::S3AuditStorage = self.into();
        s3_storage.get_proof(epoch).await
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A note to readers:
//!
//! We are marking most of the following tests as `#[ignored]` so they don't run automatically
//! with `cargo test`. These require a docker container running `minio` to be running which is
//! run in our CI pipeline when changes to this crate are made, but not otherwise since they
//! are relatively expensive integration test operations. See [s3.yml](.github/workflows/s3.yml)
//! for more information

use super::*;
use crate::common_test::AuditInformation;
use crate::storage::s3;
use crate::storage::AuditProofStorage;
use anyhow::Result;
use aws_sdk_dynamodb::types::Blob;
use aws_smithy_http::byte_stream::ByteStream;

// These are constants that are matched in both this crate's `docker-compose.yml`
// and the aws.yml workflow pipeline
const ACCESS_KEY: &str = "test";
const SECRET_KEY: &str = "someLongAccessKey";
const TEST_REGION: &str = "us-east-2";
const TEST_DYNAMO_ENDPOINT: &str = "http://127.0.0.1:9002";
const TEST_S3_ENDPOINT: &str = "http://127.0.0.1:9000";

#[test]
fn test_dynamo_table_naming() {
    let too_short = "a";
    assert!(matches!(super::validate_table_name(too_short), Err(_)));

    // take a-z until we reach 260 char's, then gen the string
    let too_long = ('a'..='z').cycle().take(260).collect::<String>();
    assert!(matches!(super::validate_table_name(&too_long), Err(_)));

    let bad_chars = "!@#$%^&*()_+";
    assert!(matches!(super::validate_table_name(bad_chars), Err(_)));

    // Ok matches
    assert!(matches!(super::validate_table_name("table123"), Ok(_)));
    assert!(matches!(
        super::validate_table_name("some-table-name"),
        Ok(_)
    ));
    assert!(matches!(
        super::validate_table_name("some.table.name"),
        Ok(_)
    ));
    assert!(matches!(
        super::validate_table_name("some_table_name"),
        Ok(_)
    ));
}

#[tokio::test]
#[ignore]
async fn integration_test_dynamo_listing() {
    // make sure we have a valid bucket name, that's "somewhat" unique
    let table = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket and table name is {}", table);
    assert!(matches!(s3::validate_bucket_name(&table), Ok(_)));
    assert!(matches!(validate_table_name(&table), Ok(_)));

    // Get the storage reader
    let storage = get_dynamo_storage(&table);
    let s3_storage: s3::S3AuditStorage = (&storage).into();

    // Note: We have 2 shared configs, because we have 2 different endpoints here which
    // we need to differentiate between. We could gin up our own endpoint resolver, but
    // this is much easier to work with
    let dynamo_shared_config = storage.get_shared_config().await;
    let s3_shared_config = s3_storage.get_shared_test_config().await;

    // Populate the test storage
    populate_test_storage(&dynamo_shared_config, &s3_shared_config, &table, 10, false)
        .await
        .expect("Failed to populate test storage");

    // List the epochs found in the storage layer
    let mut epoch_summaries: Vec<EpochSummary> = storage
        .list_proofs(ProofIndexCacheOption::NoCache)
        .await
        .expect("Failed to list proofs");
    epoch_summaries.sort_by(|a, b| a.name.epoch.cmp(&b.name.epoch));

    // There should be 10 proofs in the storage layer
    log::info!(
        "There are {} epochs in the storage layer",
        epoch_summaries.len()
    );
    assert_eq!(10, epoch_summaries.len());

    // check the linear history of the proofs
    log::info!("Checking linear history of audit proofs");
    for (i, summary) in epoch_summaries.into_iter().enumerate() {
        assert_eq!(i as u64, summary.name.epoch);
    }

    // if the test is successful, try a cleanup of the storage now
    maybe_flush_storage(&dynamo_shared_config, &s3_shared_config, &table)
        .await
        .expect("Failed to flush storage");
}

#[tokio::test]
#[ignore]
async fn integration_test_dynamo_audit_verification() {
    // make sure we have a valid bucket name, that's "somewhat" unique
    let table = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket and table name is {}", table);
    assert!(matches!(s3::validate_bucket_name(&table), Ok(_)));
    assert!(matches!(validate_table_name(&table), Ok(_)));

    // Get the storage reader
    let storage = get_dynamo_storage(&table);
    let s3_storage: s3::S3AuditStorage = (&storage).into();

    // Note: We have 2 shared configs, because we have 2 different endpoints here which
    // we need to differentiate between. We could gin up our own endpoint resolver, but
    // this is much easier to work with
    let dynamo_shared_config = storage.get_shared_config().await;
    let s3_shared_config = s3_storage.get_shared_test_config().await;

    // Populate the test storage
    populate_test_storage(&dynamo_shared_config, &s3_shared_config, &table, 3, false)
        .await
        .expect("Failed to populate test storage");

    // List the epochs found in the storage layer
    let mut epoch_summaries: Vec<EpochSummary> = storage
        .list_proofs(ProofIndexCacheOption::NoCache)
        .await
        .unwrap();
    epoch_summaries.sort_by(|a, b| a.name.epoch.cmp(&b.name.epoch));

    // There should be 3 proofs in the storage layer
    log::info!(
        "There are {} epochs in the storage layer",
        epoch_summaries.len()
    );
    assert_eq!(3, epoch_summaries.len());

    // verify all fo the audit proofs
    for epoch in epoch_summaries.iter() {
        let proof_blob = storage.get_proof(epoch).await.unwrap();
        log::info!(
            "Verification epoch {} -> {}",
            epoch.name.epoch,
            epoch.name.epoch + 1
        );
        crate::auditor::audit_epoch(proof_blob.clone(), false)
            .await
            .unwrap();
        crate::auditor::audit_epoch(proof_blob, true).await.unwrap();
    }

    // if the test is successful, try a cleanup of the storage now
    maybe_flush_storage(&dynamo_shared_config, &s3_shared_config, &table)
        .await
        .expect("Failed to flush storage");
}

#[tokio::test]
#[ignore]
async fn populate_test_dynamo() {
    // Populates the test bucket for use with the command-line REPL via the command
    // cargo run -p akd_local_auditor -- dynamo-db --table populatetestdynamo --bucket populatetestdynamo --region us-east-2 --s3-endpoint http://127.0.0.1:9000 --dynamo-endpoint http://127.0.0.1:9002 --access-key test --secret-key someLongAccessKey

    // make sure we have a valid bucket name, that's "somewhat" unique
    let table = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket and table name is {}", table);
    assert!(matches!(s3::validate_bucket_name(&table), Ok(_)));
    assert!(matches!(validate_table_name(&table), Ok(_)));

    // Get the storage reader
    let storage = get_dynamo_storage(&table);
    let s3_storage: s3::S3AuditStorage = (&storage).into();

    // Note: We have 2 shared configs, because we have 2 different endpoints here which
    // we need to differentiate between. We could gin up our own endpoint resolver, but
    // this is much easier to work with
    let dynamo_shared_config = storage.get_shared_config().await;
    let s3_shared_config = s3_storage.get_shared_test_config().await;

    // Populate the test storage
    populate_test_storage(&dynamo_shared_config, &s3_shared_config, &table, 50, false)
        .await
        .expect("Failed to populate test storage");
}

fn build_dynamo_table_properties(
    query: dynamo::client::fluent_builders::CreateTable,
) -> dynamo::client::fluent_builders::CreateTable {
    use aws_sdk_dynamodb::model::{
        AttributeDefinition, KeySchemaElement, KeyType, ProvisionedThroughput, ScalarAttributeType,
    };

    let primary_key = "epoch".to_string();

    let ad = AttributeDefinition::builder()
        .attribute_name(&primary_key)
        .attribute_type(ScalarAttributeType::N)
        .build();

    let ks = KeySchemaElement::builder()
        .attribute_name(&primary_key)
        .key_type(KeyType::Hash)
        .build();

    let pt = ProvisionedThroughput::builder()
        .read_capacity_units(100)
        .write_capacity_units(5)
        .build();

    query
        .key_schema(ks)
        .attribute_definitions(ad)
        .provisioned_throughput(pt)
}

/// Populate a test bucket with `n` audit proofs
async fn populate_test_storage(
    dynamo_shared_config: &aws_config::SdkConfig,
    s3_shared_config: &aws_config::SdkConfig,
    table: &str,
    n_blobs: usize,
    expensive: bool,
) -> Result<()> {
    // flush all the storage, so we're starting fresh
    maybe_flush_storage(dynamo_shared_config, s3_shared_config, table).await?;

    // Build the S3 config from the shared SdkConfig
    let config = aws_sdk_s3::config::Builder::from(s3_shared_config)
        .retry_config(aws_sdk_s3::RetryConfig::disabled())
        .build();
    // get the S3 & dynamo clients
    let s3_client = aws_sdk_s3::Client::from_conf(config);
    let dynamo_client = aws_sdk_dynamodb::Client::new(dynamo_shared_config);

    // create the bucket
    log::debug!("Creating S3 bucket {}", table);
    s3_client
        .create_bucket()
        .bucket(table.to_string())
        .send()
        .await?;
    log::debug!("Creating DynamoDb table {}", table);
    // create the table
    build_dynamo_table_properties(dynamo_client.create_table().table_name(table.to_string()))
        .send()
        .await?;

    log::debug!("Generating {} proofs", n_blobs);
    // Generate a block of real, verifiable audit proofs
    let proofs = crate::common_test::generate_audit_proofs(n_blobs, expensive)
        .await
        .map_err(|akd_err| anyhow::anyhow!("AKD Error generating proofs: {}", akd_err))?;

    log::info!("Uploading proofs to S3 and DynamoDb...");
    // upload each proof blob into S3 and the dynamo index
    for AuditInformation {
        chash,
        phash,
        proof,
    } in proofs
    {
        // Generate the s3 compat format
        let blobs = akd::local_auditing::generate_audit_blobs(vec![phash, chash], proof)
            .map_err(|err| anyhow::anyhow!("Error generating audit blob {:?}", err))?;
        // Grab the blob + upload it
        if let Some(blob) = blobs.first() {
            let byte_stream = ByteStream::from(blob.data.clone());
            let name = blob.name.to_string();
            s3_client
                .put_object()
                .bucket(table.to_string())
                .key(name)
                .body(byte_stream)
                .send()
                .await?;

            let request = dynamo_client
                .put_item()
                .table_name(table.to_string())
                .item("epoch", AttributeValue::N(blob.name.epoch.to_string()))
                .item(
                    "previous_hash",
                    AttributeValue::B(Blob::new(blob.name.previous_hash.to_vec())),
                )
                .item(
                    "current_hash",
                    AttributeValue::B(Blob::new(blob.name.current_hash.to_vec())),
                )
                .item("blob", AttributeValue::S(blob.name.to_string()));

            request.send().await?;
        } else {
            panic!("We should never generate an empty blob array, but if we do crash hard & fast!");
        }
    }
    Ok(())
}

async fn maybe_flush_storage(
    dynamo_shared_config: &aws_config::SdkConfig,
    s3_shared_config: &aws_config::SdkConfig,
    table: &str,
) -> Result<()> {
    // flush S3
    s3::test::maybe_flush_storage(s3_shared_config, table).await?;
    log::info!("S3 bucket {} flushed", table);

    // flush dynamo
    let client = aws_sdk_dynamodb::Client::new(dynamo_shared_config);
    // describe the table
    let table_exists = client
        .list_tables()
        .send()
        .await?
        .table_names()
        .as_ref()
        .unwrap()
        .contains(&table.into());
    if table_exists {
        // if exists, delete it
        client
            .delete_table()
            .table_name(table.to_string())
            .send()
            .await?;
    }
    log::info!("Dynamo table {} flushed", table);

    Ok(())
}

fn get_dynamo_storage(table_and_bucket_name: &str) -> DynamoDbAuditStorage {
    let clap_args = DynamoDbClapSettings {
        bucket: table_and_bucket_name.to_string(),
        region: TEST_REGION.to_string(),
        dynamo_endpoint: Some(TEST_DYNAMO_ENDPOINT.to_string()),
        table: table_and_bucket_name.to_string(),
        s3_endpoint: Some(TEST_S3_ENDPOINT.to_string()),
        access_key: Some(ACCESS_KEY.to_string()),
        secret_key: Some(SECRET_KEY.to_string()),
    };
    (&clap_args).into()
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module comprises S3 bucket READ ONLY access to download and parse
//! Audit Proofs

use super::ProofIndexCacheOption;
use anyhow::{bail, Result};
use async_trait::async_trait;
use aws_config::meta::region::RegionProviderChain;
use aws_config::RetryConfig;
use aws_sdk_s3 as s3;
use aws_smithy_http::endpoint::Endpoint;
use clap::Args;
use log::{debug, error};
use s3::output::ListObjectsV2Output;
use s3::Region;
use std::convert::TryInto;
use std::sync::Arc;
use tokio::sync::RwLock;

const MIN_BUCKET_CHARS: usize = 3;
const MAX_BUCKET_CHARS: usize = 63;
const ALLOWED_BUCKET_CHARS: [char; 38] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
    't', 'u', 'v', 'w', 'x', 'y', 'z', '.', '-', '1', '2', '3', '4', '5', '6', '7', '8', '9', '0',
];
pub const DEFAULT_AWS_REGION: &str = "us-west-2";

// These are crate-visible because the dynamo tests utilize the test functions for S3 buckets
#[cfg(test)]
pub(crate) mod test;

pub(crate) fn validate_bucket_name(s: &str) -> Result<String, String> {
    let str = s.to_string();
    if str.len() < MIN_BUCKET_CHARS || str.len() > MAX_BUCKET_CHARS {
        return Err(format!(
            "Bucket name must be between [{}, {}] characters in length. Gave {}",
            MIN_BUCKET_CHARS,
            MAX_BUCKET_CHARS,
            str.len()
        ));
    }

    for c in str.chars() {
        if !ALLOWED_BUCKET_CHARS.iter().any(|v| c == *v) {
            return Err(format!("Character '{c}' is not allowed in bucket name. Bucket names must contain lower-case letters, numbers, '-', and '.' only."));
        }
    }
    Ok(str)
}

pub(crate) fn validate_uri(s: &str) -> Result<String, String> {
    let uri: http::Uri = s.parse::<http::Uri>().map_err(|err| err.to_string())?;
    match (uri.scheme(), uri.authority()) {
        (None, None) => {
            Err("URI has no scheme or authority. Relative URIs not supported".to_string())
        }
        // if there's at least a scheme (http[s]://) or authority (www.test.com) we can construct what we need
        _ => Ok(s.to_string()),
    }
}

#[derive(Args, Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct S3ClapSettings {
    /// The S3 bucket where the audit proofs are stored
    #[clap(
        long,
        value_parser = validate_bucket_name
    )]
    pub(crate) bucket: String,

    /// The AWS region to operate in
    #[clap(long)]
    pub(crate) region: String,

    /// (OPTIONAL) An custom URI for the AWS endpoint
    #[clap(long, value_parser = validate_uri)]
    pub(crate) endpoint: Option<String>,

    /// (OPTIONAL) AWS Access key for the session
    #[clap(long)]
    pub(crate) access_key: Option<String>,

    /// (OPTIONAL) AWS secret key for the session
    #[clap(long)]
    pub(crate) secret_key: Option<String>,
}

impl super::CommonStorageClapSettings for S3ClapSettings {
    fn bucket(&self) -> String {
        self.bucket.clone()
    }

    fn region(&self) -> String {
        self.region.clone()
    }

    fn s3_endpoint(&self) -> Option<String> {
        self.endpoint.clone()
    }

    fn access_key(&self) -> Option<String> {
        self.access_key.clone()
    }

    fn secret_key(&self) -> Option<String> {
        self.secret_key.clone()
    }

    fn set_bucket(&mut self, v: String) {
        self.bucket = v;
    }
    fn set_region(&mut self, v: String) {
        self.region = v;
    }
    fn set_s3_endpoint(&mut self, v: Option<String>) {
        self.endpoint = v;
    }
    fn set_access_key(&mut self, v: Option<String>) {
        self.access_key = v;
    }
    fn set_secret_key(&mut self, v: Option<String>) {
        self.secret_key = v;
    }
}

#[derive(Debug)]
pub struct S3AuditStorage {
    /// The bucket where the audit proofs are stored
    pub(crate) bucket: String,
    /// The AWS region
    pub(crate) region: String,
    /// Customize the endpoint (useful for testing)
    pub(crate) endpoint: Option<String>,
    /// The access key
    pub(crate) access_key: Option<String>,
    /// The access secret
    pub(crate) secret_key: Option<String>,
    /// The configuration, cached for subsequent accesses
    pub(crate) config: Arc<RwLock<Option<aws_config::SdkConfig>>>,
    /// Cache of epoch summaries, populated with each call to list_blob_keys
    pub(crate) cache: Arc<RwLock<Option<Vec<super::EpochSummary>>>>,
}

impl From<&S3ClapSettings> for S3AuditStorage {
    fn from(clap: &S3ClapSettings) -> Self {
        Self {
            bucket: clap.bucket.clone(),
            region: clap.region.clone(),
            endpoint: clap.endpoint.as_ref().cloned(),
            access_key: clap.access_key.as_ref().cloned(),
            secret_key: clap.secret_key.as_ref().cloned(),
            config: Arc::new(RwLock::new(None)),
            cache: Arc::new(RwLock::new(None)),
        }
    }
}

impl S3AuditStorage {
    // exposed for test functionality only
    #[cfg(test)]
    pub async fn get_shared_test_config(&self) -> aws_types::SdkConfig {
        self.get_shared_config().await
    }

    async fn get_shared_config(&self) -> aws_types::SdkConfig {
        let mut lock = self.config.write().await;
        if let Some(config) = &*lock {
            config.clone()
        } else {
            // Get the shared AWS config
            let region_provider = RegionProviderChain::first_try(Region::new(self.region.clone()))
                .or_default_provider()
                .or_else(Region::new(DEFAULT_AWS_REGION));

            let mut shared_config_loader = aws_config::from_env().region(region_provider);

            if let (Some(access_key), Some(secret_key)) = (&self.access_key, &self.secret_key) {
                let credentials = aws_types::Credentials::from_keys(access_key, secret_key, None);
                shared_config_loader = shared_config_loader.credentials_provider(credentials);
            }

            if let Some(endpoint) = &self.endpoint {
                // THE URI IS ALREADY VALIDATED BY CLAP ARGS, hence the expect here
                let endpoint_resolver = Endpoint::immutable(endpoint.parse().expect("valid URI"));
                shared_config_loader = shared_config_loader.endpoint_resolver(endpoint_resolver);
            }

            let shared_config = shared_config_loader.load().await;

            *lock = Some(shared_config.clone());
            shared_config
        }
    }

    async fn get_config(&self) -> s3::Config {
        let shared_config = self.get_shared_config().await;
        s3::config::Builder::from(&shared_config)
            .retry_config(RetryConfig::disabled())
            .build()
    }

    async fn list_objects_v2(&self) -> s3::paginator::ListObjectsV2Paginator {
        let config = self.get_config().await;

        s3::Client::from_conf(config)
            .list_objects_v2()
            .bucket(self.bucket.clone())
            .encoding_type(s3::model::EncodingType::Url)
            .fetch_owner(false)
            .into_paginator()
    }

    async fn get_object(&self, key: &str) -> Result<s3::client::fluent_builders::GetObject> {
        // We need to retrieve the minimum version of an object, to guarantee we see the ORIGINAL
        // status of the blob. Future versions are not supported

        let attributes = s3::Client::from_conf(self.get_config().await)
            .list_object_versions()
            .bucket(self.bucket.clone())
            .prefix(key.to_string())
            .max_keys(1)
            .send()
            .await?;

        let mut version_count = 0usize;
        let mut min_time = std::time::SystemTime::now();
        let mut min_version = String::new();
        let mut first_etag = None;
        if let Some(versions) = attributes.versions() {
            for version in versions {
                if let Some(potential_key) = version.key() {
                    if potential_key == key {
                        if let (Some(mod_time), Some(version_id)) =
                            (version.last_modified(), version.version_id())
                        {
                            version_count += 1;
                            if first_etag.is_none() {
                                first_etag = Some(
                                    version
                                        .e_tag()
                                        .map(|tag| tag.to_string())
                                        .unwrap_or_else(|| "".to_string()),
                                );
                            } else {
                                let next_etag = version
                                    .e_tag()
                                    .map(|tag| tag.to_string())
                                    .unwrap_or_else(|| "".to_string());
                                if Some(next_etag.clone()) != first_etag {
                                    // etag's don't match, return error
                                    return Err(anyhow::anyhow!("There were duplicate objects with the same key that have different etags which indicates different values. This epoch cannot be trusted ({} != {})", first_etag.expect("Failed to unwrap a Some() which should be impossible here."), next_etag));
                                }
                            }

                            let this_time: std::time::SystemTime = (*mod_time).try_into()?;
                            if this_time < min_time {
                                min_version = version_id.to_string();
                                min_time = this_time;
                            }
                        }
                    }
                }
            }
        }

        if version_count == 0 {
            return Err(anyhow::anyhow!(
                "Object not found with any version information"
            ));
        }

        Ok(s3::Client::from_conf(self.get_config().await)
            .get_object()
            .bucket(self.bucket.clone())
            .version_id(min_version)
            .key(key.to_string())
            .checksum_mode(s3::model::ChecksumMode::Enabled))
    }

    fn process_streamed_response(
        part: Result<ListObjectsV2Output, s3::types::SdkError<s3::error::ListObjectsV2Error>>,
    ) -> Result<Vec<super::EpochSummary>> {
        let mut partials = vec![];
        match part {
            Err(sdk_err) => {
                return Err(anyhow::anyhow!(
                    "Error executing list_objects_v2 in S3: {}",
                    sdk_err
                ));
            }
            Ok(list_objects_v2_output) => {
                if let Some(contents) = list_objects_v2_output.contents() {
                    for object_ref in contents {
                        let key = object_ref.key().unwrap_or("MISSING_KEY");
                        let epoch_summary: Result<super::EpochSummary> = key.try_into();
                        match epoch_summary {
                            Err(parse_err) => debug!(
                                "Failed to parse {} into an EpochSummary: {}. Skipping...",
                                key, parse_err
                            ),
                            Ok(summary) => partials.push(summary),
                        }
                    }
                }
            }
        }
        Ok(partials)
    }
}

#[async_trait]
impl super::AuditProofStorage for S3AuditStorage {
    fn default_cache_control(&self) -> ProofIndexCacheOption {
        ProofIndexCacheOption::UseCache
    }

    async fn list_proofs(
        &self,
        cache_control: ProofIndexCacheOption,
    ) -> Result<Vec<super::EpochSummary>> {
        use tokio_stream::StreamExt;

        {
            if let ProofIndexCacheOption::UseCache = &cache_control {
                if let Some(cache) = &*self.cache.read().await {
                    return Ok(cache.clone());
                }
            }
        }

        let mut results = vec![];
        let client = self.list_objects_v2().await;
        let mut stream = client.page_size(1000).send();
        while let Some(result) = stream.next().await {
            let mut maybe_item = Self::process_streamed_response(result)?;
            results.append(&mut maybe_item);
        }

        if !results.is_empty() {
            let mut update = self.cache.write().await;
            *update = Some(results.clone());
        }

        Ok(results)
    }

    async fn get_proof(
        &self,
        epoch: &super::EpochSummary,
    ) -> Result<akd::local_auditing::AuditBlob> {
        let client = self.get_object(&epoch.key).await?;
        match client.send().await {
            Err(some_err) => {
                error!("Error executing get_object in S3 {}", some_err);
                bail!("Error executing get_object in S3 {}", some_err);
            }
            Ok(result) => {
                let bytes = result.body.collect().await?.into_bytes();
                Ok(akd::local_auditing::AuditBlob {
                    data: bytes.into_iter().collect::<Vec<u8>>(),
                    name: epoch.name,
                })
            }
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A note to readers:
//!
//! We are marking most of the following tests as `#[ignored]` so they don't run automatically
//! with `cargo test`. These require a docker container running `minio` to be running which is
//! run in our CI pipeline when changes to this crate are made, but not otherwise since they
//! are relatively expensive integration test operations. See [s3.yml](.github/workflows/s3.yml)
//! for more information

use crate::{
    common_test::AuditInformation,
    storage::{AuditProofStorage, EpochSummary},
};

use super::*;
use anyhow::Result;
use aws_config::SdkConfig;
use aws_sdk_s3::{
    config::Builder as S3ConfigBuilder,
    model::{BucketVersioningStatus, Delete, ObjectIdentifier, VersioningConfiguration},
    Client,
};
use aws_smithy_http::byte_stream::ByteStream;

// These are constants that are matched in both this crate's `docker-compose.yml`
// and the aws.yml workflow pipeline
const ACCESS_KEY: &str = "test";
const SECRET_KEY: &str = "someLongAccessKey";
const TEST_REGION: &str = "us-east-2";
const TEST_ENDPOINT: &str = "http://127.0.0.1:9000";

/// Flush the storage bucket of all blobs + delete the bucket (if it exists)
pub async fn maybe_flush_storage(shared_config: &SdkConfig, bucket_name: &str) -> Result<()> {
    // Build the S3 config from the shared SdkConfig
    let config = S3ConfigBuilder::from(shared_config)
        .retry_config(RetryConfig::disabled())
        .build();
    // get the S3 client
    let client = Client::from_conf(config);
    // check if the bucket exists
    let buckets = client.list_buckets().send().await?;

    if let Some(bucket_names) = buckets.buckets() {
        log::info!("Found {} buckets in the test storage", bucket_names.len());

        if bucket_names.iter().any(|bucket| {
            if let Some(maybe_bucket) = bucket.name() {
                return maybe_bucket == bucket_name;
            }
            false
        }) {
            log::info!("Found target bucket {}, deleting.", bucket_name);

            // From: https://docs.aws.amazon.com/sdk-for-rust/latest/dg/rust_s3_code_examples.html
            let objects = client.list_objects_v2().bucket(bucket_name).send().await?;

            let object_versions = client
                .list_object_versions()
                .bucket(bucket_name)
                .max_keys(1000)
                .send()
                .await?;
            let map = object_versions
                .versions()
                .unwrap()
                .iter()
                .map(|obj_ver| {
                    (
                        obj_ver.key().unwrap().to_string(),
                        obj_ver.version_id().unwrap().to_string(),
                    )
                })
                .collect::<std::collections::HashMap<String, String>>();
            let get_version = |key: &str| {
                let result = map.get(key).cloned();
                if result.is_none() {
                    panic!("Found no version for key {}", key);
                }
                result
            };

            let delete_objects = objects
                .contents()
                .unwrap_or_default()
                .iter()
                .map(|obj| {
                    ObjectIdentifier::builder()
                        .set_key(Some(obj.key().unwrap().to_string()))
                        .set_version_id(get_version(obj.key().unwrap()))
                        .build()
                })
                .collect::<Vec<_>>();

            if !delete_objects.is_empty() {
                log::info!(
                    "Bucket {} has {} objects in it which must be delete first",
                    bucket_name,
                    delete_objects.len()
                );
                client
                    .delete_objects()
                    .bucket(bucket_name)
                    .delete(Delete::builder().set_objects(Some(delete_objects)).build())
                    .send()
                    .await?;
            }
            client.delete_bucket().bucket(bucket_name).send().await?;
        }
    }
    Ok(())
}

/// Populate a test bucket with `n` audit proofs
pub async fn populate_test_storage(
    shared_config: &SdkConfig,
    bucket: &str,
    n_blobs: usize,
    expensive: bool,
) -> Result<()> {
    maybe_flush_storage(shared_config, bucket).await?;

    // Build the S3 config from the shared SdkConfig
    let config = S3ConfigBuilder::from(shared_config)
        .retry_config(RetryConfig::disabled())
        .build();
    // get the S3 client
    let client = Client::from_conf(config);
    // create the bucket
    client
        .create_bucket()
        .bucket(bucket.to_string())
        .send()
        .await?;

    // enable versioning on the bucket
    let versioning_config = VersioningConfiguration::builder()
        .set_status(Some(BucketVersioningStatus::Enabled))
        .build();
    client
        .put_bucket_versioning()
        .bucket(bucket.to_string())
        .versioning_configuration(versioning_config)
        .send()
        .await?;

    // Generate a block of real, verifiable audit proofs
    let proofs = crate::common_test::generate_audit_proofs(n_blobs, expensive)
        .await
        .map_err(|akd_err| anyhow::anyhow!("AKD Error generating proofs: {}", akd_err))?;

    // upload each proof blob into S3
    for AuditInformation {
        chash,
        phash,
        proof,
    } in proofs
    {
        // Generate the s3 compat format
        let blobs = akd::local_auditing::generate_audit_blobs(vec![phash, chash], proof)
            .map_err(|err| anyhow::anyhow!("Error generating audit blob {:?}", err))?;
        // Grab the blob + upload it
        if let Some(blob) = blobs.first() {
            let byte_stream = ByteStream::from(blob.data.clone());
            let name = blob.name.to_string();
            client
                .put_object()
                .bucket(bucket.to_string())
                .key(name)
                .body(byte_stream)
                .send()
                .await?;
        } else {
            panic!("We should never generate an empty blob array, but if we do crash hard & fast!");
        }
    }
    Ok(())
}

pub fn get_s3_storage(bucket: &str) -> S3AuditStorage {
    let clap_args = S3ClapSettings {
        bucket: bucket.to_string(),
        region: TEST_REGION.to_string(),
        endpoint: Some(TEST_ENDPOINT.to_string()),
        access_key: Some(ACCESS_KEY.to_string()),
        secret_key: Some(SECRET_KEY.to_string()),
    };
    (&clap_args).into()
}

// =========================== Test Cases =========================== //

#[tokio::test]
#[ignore]
async fn integration_test_s3_bucket_listing() {
    // make sure we have a valid bucket name, that's "somewhat" unique
    let bucket = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket name is {}", bucket);
    assert!(matches!(validate_bucket_name(&bucket), Ok(_)));

    // Get the storage reader
    let storage = get_s3_storage(&bucket);
    let shared_config = storage.get_shared_config().await;

    // Populate the test storage
    populate_test_storage(&shared_config, &bucket, 10, false)
        .await
        .expect("Failed to populate test storage");

    // List the epochs found in the storage layer
    let mut epoch_summaries: Vec<EpochSummary> = storage
        .list_proofs(ProofIndexCacheOption::NoCache)
        .await
        .unwrap();
    epoch_summaries.sort_by(|a, b| a.name.epoch.cmp(&b.name.epoch));

    // There should be 10 proofs in the storage layer
    log::info!(
        "There are {} epochs in the storage layer",
        epoch_summaries.len()
    );
    assert_eq!(10, epoch_summaries.len());

    // check the linear history of the proofs
    log::info!("Checking linear history of audit proofs");
    for (i, summary) in epoch_summaries.into_iter().enumerate() {
        assert_eq!(i as u64, summary.name.epoch);
    }

    // if the test is successful, try a cleanup of the storage now
    maybe_flush_storage(&shared_config, &bucket)
        .await
        .expect("Failed to flush storage");
}

#[tokio::test]
#[ignore]
async fn integration_test_s3_audit_verification() {
    // make sure we have a valid bucket name, that's "somewhat" unique
    let bucket = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket name is {}", bucket);
    assert!(matches!(validate_bucket_name(&bucket), Ok(_)));

    // Get the storage reader
    let storage = get_s3_storage(&bucket);
    let shared_config = storage.get_shared_config().await;

    // Populate the test storage
    populate_test_storage(&shared_config, &bucket, 3, false)
        .await
        .expect("Failed to populate test storage");

    // List the epochs found in the storage layer
    let mut epoch_summaries: Vec<EpochSummary> = storage
        .list_proofs(ProofIndexCacheOption::NoCache)
        .await
        .unwrap();
    epoch_summaries.sort_by(|a, b| a.name.epoch.cmp(&b.name.epoch));

    // There should be 3 proofs in the storage layer
    log::info!(
        "There are {} epochs in the storage layer",
        epoch_summaries.len()
    );
    assert_eq!(3, epoch_summaries.len());

    // verify all fo the audit proofs
    for epoch in epoch_summaries.iter() {
        let proof_blob = storage.get_proof(epoch).await.unwrap();
        log::info!(
            "Verification epoch {} -> {}",
            epoch.name.epoch,
            epoch.name.epoch + 1
        );
        crate::auditor::audit_epoch(proof_blob.clone(), false)
            .await
            .unwrap();
        crate::auditor::audit_epoch(proof_blob, true).await.unwrap();
    }

    // if the test is successful, try a cleanup of the storage now
    maybe_flush_storage(&shared_config, &bucket)
        .await
        .expect("Failed to flush storage");
}

// ============================ IMPORTANT ============================ //
// NEVER run this operation in a CI pipeline, it takes 40+ min to complete
// on reasonable hardware, and is only used for manual testing & timing information
#[tokio::test]
#[ignore]
async fn expensive_audit_verification() {
    // make sure we have a valid bucket name, that's "somewhat" unique
    let bucket = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket name is {}", bucket);
    assert!(matches!(validate_bucket_name(&bucket), Ok(_)));

    // Get the storage reader
    let storage = get_s3_storage(&bucket);
    let shared_config = storage.get_shared_config().await;

    // Populate the test storage
    populate_test_storage(&shared_config, &bucket, 3, true)
        .await
        .unwrap();

    // List the epochs found in the storage layer
    let mut epoch_summaries: Vec<EpochSummary> = storage
        .list_proofs(ProofIndexCacheOption::NoCache)
        .await
        .unwrap();
    epoch_summaries.sort_by(|a, b| a.name.epoch.cmp(&b.name.epoch));

    // There should be 3 proofs in the storage layer
    log::info!(
        "There are {} epochs in the storage layer",
        epoch_summaries.len()
    );
    assert_eq!(3, epoch_summaries.len());

    // Verify just 1 epoch
    let now = std::time::Instant::now();

    let epoch = &epoch_summaries[3];
    let proof_blob = storage.get_proof(epoch).await.unwrap();
    log::info!(
        "Verification epoch {} -> {}",
        epoch.name.epoch,
        epoch.name.epoch + 1
    );
    crate::auditor::audit_epoch(proof_blob.clone(), false)
        .await
        .unwrap();
    crate::auditor::audit_epoch(proof_blob, true).await.unwrap();

    // it prints '2'
    log::error!(
        "**** ELAPSED AUDITING TIME **** {}",
        now.elapsed().as_secs()
    );

    // if the test is successful, try a cleanup of the storage now
    maybe_flush_storage(&shared_config, &bucket)
        .await
        .expect("Failed to flush storage");
}

/// Populates the test bucket for use with the command-line REPL via the command
/// Populate the bucket with
/// ```bash
/// cargo test -p akd_local_auditor populate_test_bucket -- --ignored
/// ```
///
/// Connect to the bucket with
/// ```bash
/// cargo run -p akd_local_auditor -- s3 --bucket populatetestbucket --region us-east-2 --endpoint http://127.0.0.1:9000 --access-key test --secret-key someLongAccessKey
/// ```
#[tokio::test]
#[ignore]
async fn populate_test_bucket() {
    // make sure we have a valid bucket name, that's "somewhat" unique
    let bucket = crate::common_test::alphanumeric_function_name!();
    log::debug!("Test bucket name is {}", bucket);
    assert!(matches!(validate_bucket_name(&bucket), Ok(_)));

    // Get the storage reader
    let storage = get_s3_storage(&bucket);
    let shared_config = storage.get_shared_config().await;

    // Populate the test storage
    populate_test_storage(&shared_config, &bucket, 1000, false)
        .await
        .expect("Failed to populate test storage");
}

#[test]
fn test_bucket_name_parsing() {
    let too_short = "a";
    assert!(matches!(super::validate_bucket_name(too_short), Err(_)));

    let too_long = "1234567890123456789012345678901234567890123456789012345678901234567890";
    assert!(matches!(super::validate_bucket_name(too_long), Err(_)));

    let bad_chars = "!@#$%^&*()_+";
    assert!(matches!(super::validate_bucket_name(bad_chars), Err(_)));

    let ok_1 = "12345";
    let ok_2 = "some-bucket-name";
    let ok_3 = "some.bucket.name";

    assert!(matches!(super::validate_bucket_name(ok_1), Ok(_)));
    assert!(matches!(super::validate_bucket_name(ok_2), Ok(_)));
    assert!(matches!(super::validate_bucket_name(ok_3), Ok(_)));
}

#[test]
fn test_uri_parsing() {
    let non_uri = "hasdf!@#";
    let bad_port = "http://localhost:1234%567891230";

    assert!(matches!(super::validate_uri(non_uri), Err(_)));
    assert!(matches!(super::validate_uri(bad_port), Err(_)));
    assert!(matches!(super::validate_uri("/test/uri"), Err(_)));

    assert!(matches!(super::validate_uri("www.ok.com"), Ok(_)));
    assert!(matches!(super::validate_uri("http://ok.com"), Ok(_)));
    assert!(matches!(
        super::validate_uri("http://127.0.0.1:9000"),
        Ok(_)
    ));
    assert!(matches!(
        super::validate_uri("http://east-us-2.aws.com:123"),
        Ok(_)
    ));
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A user-interface for the local auditor service for easy usage

// TODO list:
// 1. Breadcrumbs
// 2. Epoch drill-down

use std::sync::Arc;

use iced::alignment;
use iced::event::{self, Event};
use iced::theme::Theme;
use iced::widget::{button, column, container, horizontal_rule, text, Row};
use iced::{keyboard, subscription, window};
use iced::{Application, Command, Element, Length, Settings, Subscription};

type Storage = Arc<Box<dyn crate::storage::AuditProofStorage>>;
use settings::StorageSettings;

mod auditor;
mod logstream;
mod settings;

pub const AUDITOR_INTERFACE_TAB: usize = 0;
pub const SETTINGS_INTERFACE_TAB: usize = 1;
pub const LOG_INTERFACE_TAB: usize = 2;

pub const DEFAULT_SPACING: u16 = 15;
pub const DEFAULT_FONT_SIZE: u16 = 20;

/*
================= Helpers =================
*/

pub trait AuditorTab {
    fn tab_title(&self) -> String;

    fn view(&self) -> Element<Message>;
}

pub fn empty_action_message(msg: Message) -> Command<Message> {
    Command::perform(async {}, |_| msg)
}

pub fn empty_message<'a>(message: &str, theme: &'a iced::Theme) -> Element<'a, Message> {
    container(
        text(message)
            .width(Length::Fill)
            .size(30)
            .horizontal_alignment(alignment::Horizontal::Center)
            .style(theme.extended_palette().secondary.weak.text),
    )
    .width(Length::Fill)
    .height(Length::Units(200))
    .center_y()
    .into()
}

/*
================= User Interface =================
*/

pub(crate) struct UserInterface {
    should_exit: bool,
    active_tab: usize,
    settings: StorageSettings,
    auditor: auditor::Auditor,
    logs: logstream::LogStream,
    theme: Theme,
}

impl UserInterface {
    pub(crate) fn execute(processor: Option<Storage>) -> anyhow::Result<()> {
        let theme = Theme::Dark;

        let mut settings = Settings::with_flags((theme, processor));
        settings.window = window::Settings {
            size: (800, 800),
            ..window::Settings::default()
        };
        Self::run(settings)
            .map_err(|err| anyhow::anyhow!("Some user-interface error occurred {}", err))
    }
}

#[derive(Debug, Clone)]
pub enum Message {
    Exit,

    Auditor(auditor::AuditorMessage),
    Settings(settings::SettingsMessage),
    SettingsUpdated,
    Logs(logstream::LoggerMessage),

    TabSelected(usize),
    UnhandledKeyboardEvent(keyboard::Event),
}
impl Application for UserInterface {
    type Message = Message;
    type Theme = Theme;
    type Executor = iced::executor::Default;
    type Flags = (Theme, Option<Storage>);

    fn new((theme, optional_settings): Self::Flags) -> (Self, Command<Message>) {
        let active_tab = if optional_settings.is_some() {
            AUDITOR_INTERFACE_TAB
        } else {
            SETTINGS_INTERFACE_TAB
        };
        (
            Self {
                auditor: auditor::Auditor::new(optional_settings, theme.clone()),
                should_exit: false,
                active_tab,
                settings: StorageSettings::new(theme.clone()),
                logs: logstream::LogStream::new(theme.clone()),
                theme,
            },
            // try and load the previous state
            Command::perform(settings::SpecificStorageSettings::load(), |r| {
                Message::Settings(settings::SettingsMessage::Loaded(r))
            }),
        )
    }

    fn title(&self) -> String {
        "AKD Auditor".to_string()
    }

    fn update(&mut self, message: Message) -> Command<Message> {
        match message {
            Message::Exit => {
                self.should_exit = true;
                Command::none()
            }
            Message::Settings(settings_msg) => self.settings.update(settings_msg),
            Message::Auditor(audit_msg) => self.auditor.update(audit_msg),
            Message::Logs(log_msg) => self.logs.update(log_msg),
            Message::TabSelected(tab_index) => {
                self.active_tab = tab_index;
                Command::none()
            }
            Message::SettingsUpdated => {
                self.auditor.settings_updated(&self.settings);
                self.active_tab = AUDITOR_INTERFACE_TAB;
                empty_action_message(Message::Auditor(auditor::AuditorMessage::RefreshEpochs))
            }
            Message::UnhandledKeyboardEvent(evt) => {
                match evt {
                    // F5 to refresh on the main auditor page
                    keyboard::Event::KeyPressed {
                        key_code: keyboard::KeyCode::F5,
                        ..
                    } if self.active_tab == AUDITOR_INTERFACE_TAB => empty_action_message(
                        Message::Auditor(auditor::AuditorMessage::RefreshEpochs),
                    ),
                    // check CTRL-S on the settings page
                    keyboard::Event::KeyPressed {
                        key_code: keyboard::KeyCode::S,
                        modifiers,
                    } if self.active_tab == SETTINGS_INTERFACE_TAB => {
                        // Returns true if a "command key" is pressed in the [`Modifiers`].
                        //
                        // The "command key" is the main modifier key used to issue commands in the
                        // current platform. Specifically:
                        //
                        // - It is the `logo` or command key (âŒ˜) on macOS
                        // - It is the `control` key on other platforms
                        if modifiers.command() {
                            empty_action_message(Message::Settings(
                                settings::SettingsMessage::SaveSettings,
                            ))
                        } else {
                            Command::none()
                        }
                    }
                    // CTRL-X exits the app
                    keyboard::Event::KeyPressed {
                        key_code: keyboard::KeyCode::X,
                        modifiers,
                    } => {
                        if modifiers.command() {
                            empty_action_message(Message::Exit)
                        } else {
                            Command::none()
                        }
                    }
                    // Tabbing around the interface should toggle focuses
                    keyboard::Event::KeyPressed {
                        key_code: keyboard::KeyCode::Tab,
                        modifiers,
                    } => {
                        if modifiers.shift() {
                            iced::widget::focus_previous()
                        } else {
                            iced::widget::focus_next()
                        }
                    }
                    _ => Command::none(),
                }
            }
        }
    }

    fn should_exit(&self) -> bool {
        self.should_exit
    }

    fn view(&self) -> Element<Message> {
        let tabs: Vec<Box<&dyn AuditorTab>> = vec![
            Box::new(&self.auditor),
            Box::new(&self.settings),
            Box::new(&self.logs),
        ];

        let mut tab_bar = Row::new();
        for (i, tab) in tabs.iter().enumerate() {
            tab_bar = tab_bar.push(
                button(text(tab.tab_title()).size(14))
                    .width(Length::Shrink)
                    .padding(DEFAULT_SPACING)
                    .style(iced::theme::Button::Secondary)
                    .on_press(Message::TabSelected(i)),
            );
            if i < tabs.len() - 1 {
                tab_bar = tab_bar.push(iced::widget::horizontal_space(Length::Fill));
            }
        }

        tab_bar = tab_bar.push(
            button(text("Exit").size(14))
                .width(Length::Shrink)
                .padding(DEFAULT_SPACING)
                .style(iced::theme::Button::Destructive)
                .on_press(Message::Exit),
        );
        tab_bar = tab_bar
            .width(Length::Fill)
            .padding(DEFAULT_SPACING)
            .spacing(DEFAULT_SPACING);

        column![
            tab_bar,
            horizontal_rule(10),
            tabs.get(self.active_tab).unwrap().view()
        ]
        .into()
    }

    fn subscription(&self) -> Subscription<Message> {
        subscription::events_with(|event, status| match (event, status) {
            (Event::Keyboard(keyboard_evt), event::Status::Ignored) => {
                Some(Message::UnhandledKeyboardEvent(keyboard_evt))
            }
            _ => None,
        })
    }

    fn theme(&self) -> Self::Theme {
        self.theme.clone()
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Breadcrumb management

// TODO:
// 1. filter to only show a few breadcrumbs (i.e. 5)

use super::super::Message;
use super::AuditorMessage;

use std::fmt::{Debug, Display};

use iced::widget::{button, container, horizontal_space, scrollable, text, Row};
use iced::{Alignment, Command, Element, Length};

const DEFAULT_FONT_SIZE: u16 = 12;
const DEFAULT_SPACING: u16 = 5;

pub struct Breadcrumbs<T: Display + Debug + Clone> {
    crumbs: Vec<T>,
    current: usize,
    theme: iced::Theme,
}

#[derive(Debug, Clone)]
pub enum BreadcrumbMessage<T> {
    Select(usize),
    Increment,
    Decrement,
    First,
    Last,
    Push(T),
}

impl<T: Display + Debug + Clone> Breadcrumbs<T> {
    pub fn new(theme: iced::Theme) -> Self {
        Self {
            crumbs: vec![],
            current: 0,
            theme,
        }
    }

    pub fn increment(&mut self) {
        if self.can_increment() {
            self.current += 1;
        } else {
            self.current = self.crumbs.len() - 1;
        }
    }

    pub fn decrement(&mut self) {
        if self.can_decrement() {
            self.current -= 1;
        } else {
            self.current = 0;
        }
    }

    pub fn goto_last(&mut self) {
        if !self.crumbs.is_empty() {
            self.current = self.crumbs.len() - 1;
        } else {
            self.current = 0;
        }
    }

    pub fn goto_first(&mut self) {
        self.current = 0;
    }

    pub fn clear(&mut self) {
        self.crumbs.clear();
        self.current = 0;
    }

    pub fn push(&mut self, t: T) {
        self.crumbs = self
            .crumbs
            .iter()
            .take(self.current + 1)
            .cloned()
            .collect::<Vec<_>>();

        self.crumbs.push(t);
        self.goto_last();
    }

    pub fn reset(&mut self, crumb: T) {
        self.clear();
        self.push(crumb);
    }

    pub fn can_increment(&self) -> bool {
        !self.crumbs.is_empty() && self.current < self.crumbs.len() - 1
    }

    pub fn can_decrement(&self) -> bool {
        !self.crumbs.is_empty() && self.current > 0
    }

    pub fn in_view(&self) -> Option<T> {
        self.crumbs.get(self.current).cloned()
    }

    pub fn view(&self) -> Element<Message> {
        let mut the_row = Row::new();
        the_row = the_row.push(
            button("First")
                .style(iced::theme::Button::Text)
                .padding(10)
                .on_press(Message::Auditor(AuditorMessage::Breadcrumb(
                    BreadcrumbMessage::First,
                ))),
        );

        // for now, a hack is to make the middle scrollable
        let scroll = scrollable(container(
            Row::with_children(
                self.crumbs
                    .iter()
                    .enumerate()
                    .map(|(crumb_id, crumb)| {
                        let mut crumb_text = text(crumb.to_string()).size(DEFAULT_FONT_SIZE);
                        if crumb_id == self.current {
                            // color the selected item RED
                            crumb_text = crumb_text.style(self.theme.palette().primary);
                        } else {
                            crumb_text = crumb_text.style(self.theme.palette().text);
                        }
                        let mut crumb_set = Row::with_children(vec![button(crumb_text)
                            .style(iced::theme::Button::Text)
                            .padding(DEFAULT_SPACING)
                            .on_press(Message::Auditor(AuditorMessage::Breadcrumb(
                                BreadcrumbMessage::Select(crumb_id),
                            )))
                            .into()]);
                        if crumb_id != self.crumbs.len() - 1 {
                            // use a no-op button to just be the same size so we get proper spacing between the ">"s
                            crumb_set = crumb_set.push(
                                button(text(">").size(DEFAULT_FONT_SIZE))
                                    .style(iced::theme::Button::Text)
                                    .padding(DEFAULT_SPACING),
                            );
                        }
                        crumb_set.into()
                    })
                    .collect(),
            )
            .width(Length::Fill)
            .padding(DEFAULT_SPACING)
            .height(Length::Fill),
        ));
        the_row = the_row.align_items(Alignment::Center).push(scroll);

        the_row = the_row.push(horizontal_space(Length::Fill));
        the_row = the_row.push(
            button("Last")
                .style(iced::theme::Button::Text)
                .padding(DEFAULT_SPACING)
                .on_press(Message::Auditor(AuditorMessage::Breadcrumb(
                    BreadcrumbMessage::Last,
                ))),
        );

        the_row.width(Length::Fill).into()
    }

    pub fn update(&mut self, message: BreadcrumbMessage<T>) -> Command<Message> {
        match message {
            BreadcrumbMessage::Select(crumb_id) => {
                self.current = crumb_id;
            }
            BreadcrumbMessage::Increment => self.increment(),
            BreadcrumbMessage::Decrement => self.decrement(),
            BreadcrumbMessage::First => {
                self.goto_first();
            }
            BreadcrumbMessage::Last => {
                self.goto_last();
            }
            BreadcrumbMessage::Push(item) => {
                self.push(item);
            }
        }
        Command::none()
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! The main interface for doing audit operations

use self::breadcrumb::{BreadcrumbMessage, Breadcrumbs};

use super::{empty_message, Message, Storage, DEFAULT_FONT_SIZE, DEFAULT_SPACING};
use crate::storage;
use crate::storage::EpochSummary;

use iced::alignment;
use iced::widget::{
    button, column, container, horizontal_rule, horizontal_space, qr_code, row, scrollable, text,
    Column, Row,
};
use iced::{Alignment, Command, Element, Length};
use std::collections::HashMap;

mod breadcrumb;

const CRUMB_SUMMARY_CUTOFF: usize = 10;
const CRUMB_SUMMARY_COLUMNS: usize = 3;

enum ProcessingMode {
    None,
    Refreshing,
    Auditing,
}

pub struct Auditor {
    proofs: Vec<EpochSummary>,
    storage: Option<Storage>,
    qr_codes: HashMap<EpochSummary, Result<qr_code::State, String>>,
    crumbs: Breadcrumbs<BreadcrumbValue>,
    processing_mode: ProcessingMode,
    theme: iced::Theme,
}

#[derive(Clone, Debug)]
pub struct BreadcrumbValue {
    pub low: u64,
    pub high: u64,
}

impl std::fmt::Display for BreadcrumbValue {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}..{}", self.low, self.high)
    }
}

#[derive(Debug, Clone)]
pub enum AuditorMessage {
    Breadcrumb(BreadcrumbMessage<BreadcrumbValue>),

    RefreshEpochs,
    RefreshComplete(Result<Vec<EpochSummary>, String>),

    AuditEpoch(EpochSummary),
    AuditComplete((EpochSummary, Result<Vec<u8>, String>)),
}

impl From<AuditorMessage> for Message {
    fn from(a: AuditorMessage) -> Self {
        Message::Auditor(a)
    }
}

impl Auditor {
    pub fn settings_updated(&mut self, settings: &super::settings::StorageSettings) {
        let storage: Box<dyn storage::AuditProofStorage> = match &settings.settings {
            super::settings::SpecificStorageSettings::S3(s3_settings) => {
                let imp: storage::s3::S3AuditStorage = s3_settings.into();
                Box::new(imp)
            }
            super::settings::SpecificStorageSettings::DynamoDb(dynamo_settings) => {
                let imp: storage::dynamodb::DynamoDbAuditStorage = dynamo_settings.into();
                Box::new(imp)
            }
        };

        self.storage = Some(std::sync::Arc::new(storage));
    }

    pub fn new(storage: Option<Storage>, theme: iced::Theme) -> Self {
        Self {
            storage,
            crumbs: Breadcrumbs::new(theme.clone()),
            proofs: vec![],
            // qr: None,
            qr_codes: HashMap::new(),
            processing_mode: ProcessingMode::None,
            theme,
        }
    }

    async fn refresh_epochs(maybe_storage: Option<Storage>) -> Result<Vec<EpochSummary>, String> {
        if let Some(storage) = maybe_storage {
            storage
                .list_proofs(crate::storage::ProofIndexCacheOption::NoCache)
                .await
                .map_err(|e| e.to_string())
        } else {
            Err("No storage is available!".to_string())
        }
    }

    async fn audit_epoch(
        maybe_storage: Option<Storage>,
        epoch_summary: EpochSummary,
    ) -> Result<Vec<u8>, String> {
        if let Some(storage) = maybe_storage {
            // download the proof
            let blob = storage
                .get_proof(&epoch_summary)
                .await
                .map_err(|err| format!("Error downloading proof {err}"))?;
            // decode the proof
            let (epoch, p_hash, c_hash, proof) = blob
                .decode()
                .map_err(|err| format!("Error decodeing proof {err:?}"))?;
            // audit the proof
            akd::auditor::audit_verify(
                vec![p_hash, c_hash],
                akd::AppendOnlyProof {
                    proofs: vec![proof],
                    epochs: vec![epoch],
                },
            )
            .await
            .map_err(|err| format!("Error auditing proof {err}"))?;

            let qr_data = crate::auditor::format_qr_record(p_hash, c_hash, epoch);
            Ok(qr_data)
        } else {
            Err("No storage is available!".to_string())
        }
    }

    pub fn update(&mut self, message: AuditorMessage) -> Command<Message> {
        match message {
            AuditorMessage::RefreshEpochs => {
                self.processing_mode = ProcessingMode::Refreshing;

                Command::perform(Self::refresh_epochs(self.storage.clone()), |args| {
                    Message::Auditor(AuditorMessage::RefreshComplete(args))
                })
            }
            AuditorMessage::RefreshComplete(data) => {
                self.processing_mode = ProcessingMode::None;

                match data {
                    Ok(mut epochs) => {
                        log::info!("Received {} epochs", epochs.len());
                        epochs.sort_unstable_by(|a, b| a.name.epoch.cmp(&b.name.epoch));

                        // reset the crumbs
                        let crumb = BreadcrumbValue {
                            low: epochs.first().unwrap().name.epoch,
                            high: epochs.last().unwrap().name.epoch,
                        };
                        self.crumbs.reset(crumb);

                        self.proofs = epochs;
                        super::logstream::info(format!("Received {} epochs", self.proofs.len()))
                    }
                    Err(err) => {
                        let c = super::logstream::error(format!(
                            "Error downloading epoch information {err}"
                        ));
                        log::error!("Error downloading epoch infomation {}", err);
                        c
                    }
                }
            }

            AuditorMessage::AuditEpoch(epoch) => {
                self.processing_mode = ProcessingMode::Auditing;
                let epoch_clone = epoch.clone();
                Command::perform(
                    Self::audit_epoch(self.storage.clone(), epoch),
                    |args: Result<Vec<u8>, String>| {
                        Message::Auditor(AuditorMessage::AuditComplete((epoch_clone, args)))
                    },
                )
            }
            AuditorMessage::AuditComplete((epoch_summary, result)) => {
                self.processing_mode = ProcessingMode::None;
                match result {
                    Ok(msg) => {
                        let command = super::logstream::info(format!(
                            "Audit succeeded for epoch {}. Generating qr code...",
                            epoch_summary.name.epoch
                        ));
                        log::info!("Audit succeeded! Generating QR code...");
                        if let Ok(state) = qr_code::State::new(msg) {
                            self.qr_codes.insert(epoch_summary, Ok(state));
                            command
                        } else {
                            log::warn!("Failed to generate auditor verification QR code");
                            Command::batch([
                                command,
                                super::logstream::warn(format!("Failed to generate auditor verification QR code, but audit was successful for epoch {}", epoch_summary.name.epoch))
                            ])
                        }
                    }
                    Err(err) => {
                        let command = super::logstream::error(format!(
                            "Audit of epoch {} failed with error {}",
                            epoch_summary.name.epoch, err
                        ));
                        log::error!(
                            "Audit of epoch {} failed with error {}",
                            epoch_summary.name.epoch,
                            err
                        );

                        self.qr_codes
                            .insert(epoch_summary, Err(format!("Audit failed! {err}")));
                        command
                    }
                }
            }

            AuditorMessage::Breadcrumb(crumb_msg) => self.crumbs.update(crumb_msg),
        }
    }

    fn is_summary_view(&self) -> bool {
        if let Some(current) = self.crumbs.in_view() {
            current.high - current.low > (CRUMB_SUMMARY_CUTOFF as u64)
        } else {
            false
        }
    }

    fn generate_summary_view(&self) -> Element<Message> {
        let current_range = self.crumbs.in_view().unwrap();
        let chunk_size = ((current_range.high - current_range.low + 1) as f64
            / (CRUMB_SUMMARY_CUTOFF as f64).floor()) as usize;
        let summaries = (current_range.low..=current_range.high)
            .collect::<Vec<_>>()
            .chunks(chunk_size)
            .map(|chunk| (*chunk.first().unwrap(), *chunk.last().unwrap()))
            .collect::<Vec<_>>();

        // Create CRUMB_SUMMARY_COLUMNS column layouts (rows), each with
        // an equal number of row layouts (columns)
        Column::with_children(
            summaries
                .chunks(CRUMB_SUMMARY_COLUMNS)
                .map(|row_batch| {
                    Row::with_children(
                        row_batch
                            .iter()
                            .map(|(start, end)| {
                                let txt = text(format!("{start}..{end}")).size(DEFAULT_FONT_SIZE);
                                button(txt)
                                    .padding(DEFAULT_SPACING)
                                    .width(Length::FillPortion(1))
                                    .on_press(Message::Auditor(AuditorMessage::Breadcrumb(
                                        BreadcrumbMessage::Push(BreadcrumbValue {
                                            low: *start,
                                            high: *end,
                                        }),
                                    )))
                                    .style(iced::theme::Button::Secondary)
                                    .into()
                            })
                            .collect(),
                    )
                    .align_items(Alignment::Center)
                    .padding(DEFAULT_SPACING)
                    .spacing(DEFAULT_SPACING)
                    .into()
                })
                .collect(),
        )
        .into()
    }

    fn generate_regular_view(&self) -> Element<Message> {
        match self.processing_mode {
            ProcessingMode::Auditing => {
                return empty_message("Processing audit...", &self.theme);
            }
            ProcessingMode::Refreshing => {
                return empty_message("Refreshing epochs...", &self.theme);
            }
            _ => {}
        }

        if let Some(BreadcrumbValue { low, high }) = self.crumbs.in_view() {
            let skip = low as usize;
            let len = (high - low + 1) as usize;
            let filtered_proofs = self
                .proofs
                .iter()
                .skip(skip)
                .take(len)
                .cloned()
                .collect::<Vec<_>>();

            let epochs = Column::with_children(
                filtered_proofs
                    .iter()
                    .map(|proof| {
                        let mut data_row = row![
                            text(format!(
                                "Epoch {} -> {}:",
                                proof.name.epoch,
                                proof.name.epoch + 1
                            ))
                            .style(self.theme.palette().primary),
                            button("Verify")
                                .width(Length::Shrink)
                                .padding(DEFAULT_SPACING)
                                .on_press(Message::Auditor(AuditorMessage::AuditEpoch(
                                    proof.clone()
                                )))
                                .style(iced::theme::Button::Primary),
                            horizontal_space(Length::Fill),
                        ];

                        match self.qr_codes.get(proof) {
                            Some(Ok(qr_state)) => {
                                let qr = qr_code::QRCode::new(qr_state).cell_size(2);
                                data_row = data_row.push(qr);
                            }
                            Some(Err(audit_failure)) => {
                                let msg = text(audit_failure)
                                    .size(16)
                                    .horizontal_alignment(alignment::Horizontal::Center);
                                data_row = data_row.push(msg);
                            }
                            _ => {}
                        }

                        data_row
                            .padding(DEFAULT_SPACING)
                            .spacing(DEFAULT_SPACING)
                            .width(Length::Fill)
                            .align_items(Alignment::Center)
                    })
                    .map(Element::from)
                    .collect(),
            );

            scrollable(
                container(epochs)
                    .width(Length::Fill)
                    .padding(DEFAULT_SPACING)
                    .center_x(),
            )
            .into()
        } else {
            column![].into()
        }
    }

    fn generate_view_header(&self) -> Element<Message> {
        row![
            button(text("Refresh epochs").size(16))
                .width(Length::Shrink)
                .on_press(Message::Auditor(AuditorMessage::RefreshEpochs)),
            self.crumbs.view(),
        ]
        .align_items(alignment::Alignment::Center)
        .width(Length::Fill)
        .into()
    }
}

impl super::AuditorTab for Auditor {
    fn tab_title(&self) -> String {
        "Auditor".to_string()
    }

    fn view(&self) -> Element<Message> {
        let content = if self.is_summary_view() {
            self.generate_summary_view()
        } else {
            self.generate_regular_view()
        };

        column![self.generate_view_header(), horizontal_rule(20), content,]
            .spacing(DEFAULT_SPACING)
            .width(Length::Fill)
            .height(Length::Fill)
            .into()
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A log-stream interface for the auditor process

use super::{AuditorTab, Message, DEFAULT_SPACING};

use chrono::prelude::*;
use iced::widget::{button, column, container, scrollable, text, Column};
use iced::{Command, Element, Length};

#[derive(Debug, Clone)]
pub struct LogEntry {
    level: log::Level,
    time: DateTime<Local>,
    message: String,
}

impl std::fmt::Display for LogEntry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "[{}] {}  {}",
            self.time,
            self.level.as_str(),
            self.message
        )
    }
}

impl LogEntry {
    fn get_color(&self, theme: &iced::Theme) -> iced::Color {
        let palette = theme.palette();
        match self.level {
            log::Level::Error => palette.danger,
            log::Level::Warn => palette.primary,
            log::Level::Info => palette.success,
            _ => palette.text,
        }
    }
}

/* ======== Logger call utility functions ======== */
pub fn log(message: String, level: log::Level) -> Command<Message> {
    super::empty_action_message(Message::Logs(LoggerMessage::Push(LogEntry {
        message,
        time: Local::now(),
        level,
    })))
}
pub fn error(message: String) -> Command<Message> {
    log(message, log::Level::Error)
}
pub fn warn(message: String) -> Command<Message> {
    log(message, log::Level::Warn)
}
pub fn info(message: String) -> Command<Message> {
    log(message, log::Level::Info)
}
pub fn debug(message: String) -> Command<Message> {
    log(message, log::Level::Debug)
}

pub struct LogStream {
    messages: Vec<LogEntry>,
    theme: iced::Theme,
}

#[derive(Debug, Clone)]
pub enum LoggerMessage {
    Push(LogEntry),
    Clear,
}

impl LogStream {
    pub fn new(theme: iced::Theme) -> Self {
        Self {
            messages: vec![],
            theme,
        }
    }

    pub fn push(&mut self, msg: LogEntry) {
        self.messages.push(msg)
    }

    pub fn clear(&mut self) {
        self.messages.clear();
    }
}

impl LogStream {
    pub fn update(&mut self, msg: LoggerMessage) -> Command<Message> {
        match msg {
            LoggerMessage::Push(s) => self.push(s),
            LoggerMessage::Clear => self.clear(),
        }
        Command::none()
    }
}

impl AuditorTab for LogStream {
    fn tab_title(&self) -> String {
        "Auditor Logs".to_string()
    }

    fn view(&self) -> Element<Message> {
        let data_column = Column::with_children(
            self.messages
                .iter()
                .map(|msg| {
                    text(msg.to_string())
                        .size(16)
                        .style(msg.get_color(&self.theme))
                })
                .map(Element::from)
                .collect(),
        );
        column![
            scrollable(
                container(data_column)
                    .width(Length::Fill)
                    .padding(DEFAULT_SPACING),
            ),
            button("Clear").on_press(Message::Logs(LoggerMessage::Clear))
        ]
        .into()
    }
}