// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#![allow(clippy::let_unit_value)]

use anyhow::{Context, Error};
use async_trait::async_trait;
use fidl::endpoints::{ClientEnd, DiscoverableProtocolMarker, Proxy, ServerEnd};
use fidl_fuchsia_io::DirectoryProxy;
use fidl_fuchsia_paver::PaverRequestStream;
use fidl_fuchsia_pkg_ext::RepositoryConfigs;
use fuchsia_component::server::ServiceFs;
use fuchsia_component_test::LocalComponentHandles;
use fuchsia_merkle::Hash;
use fuchsia_pkg_testing::serve::ServedRepository;
use fuchsia_pkg_testing::{Package, RepositoryBuilder};
use fuchsia_sync::Mutex;
use futures::prelude::*;
use isolated_ota::{OmahaConfig, UpdateUrlSource};
use mock_omaha_server::{
    OmahaResponse, OmahaServer, OmahaServerBuilder, ResponseAndMetadata, ResponseMap,
};
use mock_paver::{MockPaverService, MockPaverServiceBuilder};
use std::collections::{BTreeMap, BTreeSet};
use std::io::Write;
use std::str::FromStr;
use std::sync::Arc;
use tempfile::TempDir;
use {fidl_fuchsia_io as fio, fuchsia_async as fasync};

pub const GLOBAL_SSL_CERTS_PATH: &str = "/config/ssl";
const EMPTY_REPO_PATH: &str = "/pkg/empty-repo";
const TEST_CERTS_PATH: &str = "/pkg/data/ssl";
const TEST_REPO_URL: &str = "fuchsia-pkg://integration.test.fuchsia.com";

pub enum OmahaState {
    /// Don't use Omaha for this update, instead use the provided, or default, update package URL.
    Disabled(Option<fuchsia_url::AbsolutePackageUrl>),
    /// Set up an Omaha server automatically.
    Auto(OmahaResponse),
    /// Pass the given OmahaConfig to Omaha.
    Manual(OmahaConfig),
}

pub struct TestParams {
    pub blobfs: Option<ClientEnd<fio::DirectoryMarker>>,
    pub board: String,
    pub channel: String,
    pub expected_blobfs_contents: BTreeSet<Hash>,
    pub paver: Arc<MockPaverService>,
    pub repo_config_dir: TempDir,
    pub ssl_certs: DirectoryProxy,
    pub update_merkle: Hash,
    pub version: String,
    pub update_url_source: UpdateUrlSource,
    pub paver_connector: ClientEnd<fio::DirectoryMarker>,
}

/// Connects the local component to a mock paver.
///
/// Unlike other mocks, the `fuchsia.paver.Paver` is serviced by [`isolated_ota_env::TestEnv`], so
/// this function proxies to the given `paver_dir_proxy` which is expected to host a
/// file named "fuchsia.paver.Paver" which implements the `fuchsia.paver.Paver` FIDL protocol.
pub async fn expose_mock_paver(
    handles: LocalComponentHandles,
    paver_dir_proxy: fio::DirectoryProxy,
) -> Result<(), Error> {
    let mut fs = ServiceFs::new();

    fs.dir("svc").add_service_connector(
        move |server_end: ServerEnd<fidl_fuchsia_paver::PaverMarker>| {
            fdio::service_connect_at(
                paver_dir_proxy.as_channel().as_ref(),
                &format!("/{}", fidl_fuchsia_paver::PaverMarker::PROTOCOL_NAME),
                server_end.into_channel(),
            )
            .expect("failed to connect to paver service node");
        },
    );

    fs.serve_connection(handles.outgoing_dir).expect("failed to serve paver fs connection");
    fs.collect::<()>().await;
    Ok(())
}

#[async_trait(?Send)]
pub trait TestExecutor<R> {
    async fn run(&self, params: TestParams) -> R;
}

pub struct TestEnvBuilder<R> {
    blobfs: Option<ClientEnd<fio::DirectoryMarker>>,
    board: String,
    channel: String,
    omaha: OmahaState,
    packages: Vec<Package>,
    paver: MockPaverServiceBuilder,
    repo_config: Option<RepositoryConfigs>,
    version: String,
    test_executor: Option<Box<dyn TestExecutor<R>>>,
    // The zbi and optional vbmeta contents.
    fuchsia_image: Option<(Vec<u8>, Option<Vec<u8>>)>,
    // The zbi and optional vbmeta contents of the recovery partition.
    recovery_image: Option<(Vec<u8>, Option<Vec<u8>>)>,
    firmware_images: BTreeMap<String, Vec<u8>>,
}

impl<R> TestEnvBuilder<R> {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        TestEnvBuilder {
            blobfs: None,
            board: "test-board".to_owned(),
            channel: "test".to_owned(),
            omaha: OmahaState::Disabled(Some(fuchsia_url::AbsolutePackageUrl::new(
                TEST_REPO_URL.parse().unwrap(),
                "update".parse().unwrap(),
                None,
                None,
            ))),
            packages: vec![],
            paver: MockPaverServiceBuilder::new(),
            repo_config: None,
            version: "0.1.2.3".to_owned(),
            test_executor: None,
            fuchsia_image: None,
            recovery_image: None,
            firmware_images: BTreeMap::new(),
        }
    }

    /// Add a package to the repository generated by this TestEnvBuilder.
    /// The package will also be listed in the generated update package
    /// so that it will be downloaded as part of the OTA.
    pub fn add_package(mut self, pkg: Package) -> Self {
        self.packages.push(pkg);
        self
    }

    pub fn blobfs(mut self, client: ClientEnd<fio::DirectoryMarker>) -> Self {
        self.blobfs = Some(client);
        self
    }

    /// Provide a TUF repository configuration to the package resolver.
    /// This will override the repository that the builder would otherwise generate.
    pub fn repo_config(mut self, repo: RepositoryConfigs) -> Self {
        self.repo_config = Some(repo);
        self
    }

    /// Enable/disable Omaha. OmahaState::Auto will automatically set up an Omaha server and tell
    /// the updater to use it.
    pub fn omaha_state(mut self, state: OmahaState) -> Self {
        self.omaha = state;
        self
    }

    /// Mutate the MockPaverServiecBuilder used by this TestEnvBuilder.
    pub fn paver<F>(mut self, func: F) -> Self
    where
        F: FnOnce(MockPaverServiceBuilder) -> MockPaverServiceBuilder,
    {
        self.paver = func(self.paver);
        self
    }

    pub fn test_executor(mut self, executor: Box<dyn TestExecutor<R>>) -> Self {
        self.test_executor = Some(executor);
        self
    }

    /// The zbi and optional vbmeta images to write.
    pub fn fuchsia_image(mut self, zbi: Vec<u8>, vbmeta: Option<Vec<u8>>) -> Self {
        assert_eq!(self.fuchsia_image, None);
        self.fuchsia_image = Some((zbi, vbmeta));
        self
    }

    /// The zbi and optional vbmeta images to write to the recovery partition.
    pub fn recovery_image(mut self, zbi: Vec<u8>, vbmeta: Option<Vec<u8>>) -> Self {
        assert_eq!(self.recovery_image, None);
        self.recovery_image = Some((zbi, vbmeta));
        self
    }

    /// A firmware image to write.
    pub fn firmware_image(mut self, type_: String, content: Vec<u8>) -> Self {
        assert_eq!(self.firmware_images.insert(type_, content), None);
        self
    }

    /// Turn this |TestEnvBuilder| into a |TestEnv|
    pub async fn build(mut self) -> Result<TestEnv<R>, Error> {
        let (repo_config, served_repo, ssl_certs, expected_blobfs_contents, merkle) = if self
            .repo_config
            .is_none()
        {
            // If no repo config was specified, host a repo containing the provided packages,
            // and an update package containing given images + all packages in the repo.
            let mut update =
                fuchsia_pkg_testing::UpdatePackageBuilder::new(TEST_REPO_URL.parse().unwrap())
                    .packages(
                        self.packages
                            .iter()
                            .map(|p| {
                                fuchsia_url::PinnedAbsolutePackageUrl::new(
                                    TEST_REPO_URL.parse().unwrap(),
                                    p.name().clone(),
                                    None,
                                    *p.hash(),
                                )
                            })
                            .collect::<Vec<_>>(),
                    )
                    .firmware_images(self.firmware_images);
            if let Some((zbi, vbmeta)) = self.fuchsia_image {
                update = update.fuchsia_image(zbi, vbmeta);
            }
            if let Some((zbi, vbmeta)) = self.recovery_image {
                update = update.recovery_image(zbi, vbmeta);
            }
            let (update, images) = update.build().await;

            // Do not include the images package, system-updater triggers GC after resolving it.
            let expected_blobfs_contents = self
                .packages
                .iter()
                .chain([update.as_package()])
                .flat_map(|p| p.list_blobs())
                .collect();

            let repo = Arc::new(
                self.packages
                    .iter()
                    .chain([update.as_package(), &images])
                    .fold(
                        RepositoryBuilder::from_template_dir(EMPTY_REPO_PATH)
                            .add_package(update.as_package()),
                        |repo, package| repo.add_package(package),
                    )
                    .build()
                    .await
                    .expect("build repo"),
            );

            let served_repo = Arc::clone(&repo).server().start().expect("serve repo");
            let config = RepositoryConfigs::Version1(vec![
                served_repo.make_repo_config(TEST_REPO_URL.parse().expect("make repo config"))
            ]);

            let update_merkle = *update.as_package().hash();
            // Add the update package to the list of packages, so that TestResult::check_packages
            // will expect to see the update package's blobs in blobfs.
            let mut packages = vec![update.into_package()];
            packages.append(&mut self.packages);
            (
                config,
                Some(served_repo),
                fuchsia_fs::directory::open_in_namespace(
                    TEST_CERTS_PATH,
                    fio::OpenFlags::RIGHT_READABLE,
                )
                .unwrap(),
                expected_blobfs_contents,
                update_merkle,
            )
        } else {
            // Use the provided repo config. Assume that this means we'll actually want to use
            // real SSL certificates, and that we don't need to host our own repository.
            (
                self.repo_config.unwrap(),
                None,
                fuchsia_fs::directory::open_in_namespace(
                    GLOBAL_SSL_CERTS_PATH,
                    fio::OpenFlags::RIGHT_READABLE,
                )
                .unwrap(),
                BTreeSet::new(),
                Hash::from_str("deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef")
                    .expect("make merkle"),
            )
        };

        let dir = tempfile::tempdir()?;
        let mut path = dir.path().to_owned();
        path.push("repo_config.json");
        let path = path.as_path();
        let mut file =
            std::io::BufWriter::new(std::fs::File::create(path).context("creating file")?);
        serde_json::to_writer(&mut file, &repo_config).unwrap();
        file.flush().unwrap();

        Ok(TestEnv {
            blobfs: self.blobfs,
            board: self.board,
            channel: self.channel,
            omaha: self.omaha,
            expected_blobfs_contents,
            paver: Arc::new(self.paver.build()),
            _repo: served_repo,
            repo_config_dir: dir,
            ssl_certs,
            update_merkle: merkle,
            version: self.version,
            test_executor: self.test_executor.expect("test executor must be set"),
        })
    }
}

pub struct TestEnv<R> {
    blobfs: Option<ClientEnd<fio::DirectoryMarker>>,
    channel: String,
    omaha: OmahaState,
    expected_blobfs_contents: BTreeSet<Hash>,
    paver: Arc<MockPaverService>,
    _repo: Option<ServedRepository>,
    repo_config_dir: tempfile::TempDir,
    ssl_certs: DirectoryProxy,
    update_merkle: Hash,
    board: String,
    version: String,
    test_executor: Box<dyn TestExecutor<R>>,
}

impl<R> TestEnv<R> {
    fn start_omaha(omaha: OmahaState, merkle: Hash) -> Result<UpdateUrlSource, Error> {
        match omaha {
            OmahaState::Disabled(url) => Ok(match url {
                Some(url) => UpdateUrlSource::UpdateUrl(url),
                None => UpdateUrlSource::UseDefault,
            }),
            OmahaState::Manual(cfg) => Ok(UpdateUrlSource::OmahaConfig(cfg)),
            OmahaState::Auto(response) => {
                let server = OmahaServerBuilder::default()
                    .responses_by_appid(
                        vec![(
                            "integration-test-appid".to_string(),
                            ResponseAndMetadata { response, merkle, ..Default::default() },
                        )]
                        .into_iter()
                        .collect::<ResponseMap>(),
                    )
                    .build()
                    .unwrap();
                let addr = OmahaServer::start(Arc::new(Mutex::new(server)))
                    .context("Starting omaha server")?;
                let config =
                    OmahaConfig { app_id: "integration-test-appid".to_owned(), server_url: addr };

                Ok(UpdateUrlSource::OmahaConfig(config))
            }
        }
    }

    /// Run the update, consuming this |TestEnv| and returning a |TestResult|.
    pub async fn run(self) -> R {
        let update_url_source = TestEnv::<R>::start_omaha(self.omaha, self.update_merkle)
            .expect("Starting Omaha server");

        let mut service_fs = ServiceFs::new();
        let paver_clone = Arc::clone(&self.paver);
        service_fs.add_fidl_service(move |stream: PaverRequestStream| {
            fasync::Task::spawn(
                Arc::clone(&paver_clone)
                    .run_paver_service(stream)
                    .unwrap_or_else(|e| panic!("Failed to run mock paver: {e:?}")),
            )
            .detach();
        });

        let (client, server) =
            fidl::endpoints::create_endpoints::<fidl_fuchsia_io::DirectoryMarker>();
        service_fs
            .serve_connection(server.into_channel().into())
            .expect("Failed to serve connection");
        fasync::Task::spawn(service_fs.collect()).detach();

        let params = TestParams {
            blobfs: self.blobfs,
            board: self.board,
            channel: self.channel,
            expected_blobfs_contents: self.expected_blobfs_contents,
            paver: self.paver,
            repo_config_dir: self.repo_config_dir,
            ssl_certs: self.ssl_certs,
            update_merkle: self.update_merkle,
            version: self.version,
            update_url_source,
            paver_connector: client,
        };

        self.test_executor.run(params).await
    }
}