// Copyright (c) 2019-2022 Alibaba Cloud
// Copyright (c) 2019-2022 Ant Group
//
// SPDX-License-Identifier: Apache-2.0
//

use super::{Rootfs, ROOTFS};
use crate::{
    block_device::agent_storage_source_from_block_config,
    share_fs::{do_get_guest_path, do_get_host_path},
};
use agent::Storage;
use anyhow::{anyhow, Context, Result};
use async_trait::async_trait;
use hypervisor::{
    BlockConfigModern, BlockDeviceAio, device::{
        DeviceConfig, DeviceType, device_manager::{DeviceManager, do_handle_device, get_block_device_info},
    },
};
use kata_types::config::hypervisor::VIRTIO_PMEM;
use kata_types::fs::VM_ROOTFS_FILESYSTEM_XFS;
use kata_types::mount::Mount;
use nix::sys::stat::{self, SFlag};
use oci_spec::runtime as oci;
use std::fs;
use tokio::sync::RwLock;

const BLOCKFILE_ROOTFS_FLAG: &str = "loop";

pub(crate) struct BlockRootfs {
    guest_path: String,
    device_id: String,
    mount: oci::Mount,
    storage: Option<agent::Storage>,
}

impl BlockRootfs {
    pub async fn new(
        d: &RwLock<DeviceManager>,
        sid: &str,
        cid: &str,
        dev_id: u64,
        rootfs: &Mount,
    ) -> Result<Self> {
        let container_path = do_get_guest_path(ROOTFS, cid, false, false);
        let host_path = do_get_host_path(ROOTFS, sid, cid, false, false);
        // Create rootfs dir on host to make sure mount point in guest exists, as readonly dir is
        // shared to guest via virtiofs, and guest is unable to create rootfs dir.
        fs::create_dir_all(&host_path)
            .map_err(|e| anyhow!("failed to create rootfs dir {}: {:?}", host_path, e))?;

        let blkdev_info = get_block_device_info(d).await;
        let block_driver = blkdev_info.block_device_driver.clone();
        let block_device_config = &mut BlockConfigModern {
            major: stat::major(dev_id) as i64,
            minor: stat::minor(dev_id) as i64,
            driver_option: block_driver.clone(),
            path_on_host: rootfs.source.clone(),
            blkdev_aio: BlockDeviceAio::new(&blkdev_info.block_device_aio),
            num_queues: blkdev_info.num_queues,
            queue_size: blkdev_info.queue_size,
            logical_sector_size: blkdev_info.block_device_logical_sector_size,
            physical_sector_size: blkdev_info.block_device_physical_sector_size,
            ..Default::default()
        };

        // create and insert block device into Kata VM
        let device_info = do_handle_device(d, &DeviceConfig::BlockCfgModern(block_device_config.clone()))
            .await
            .context("do handle device failed.")?;

        let mut storage = Storage {
            fs_type: rootfs.fs_type.clone(),
            mount_point: container_path.clone(),
            options: rootfs.options.clone(),
            ..Default::default()
        };

        // XFS rootfs: add 'nouuid' to avoid UUID conflicts when the same
        // disk image is mounted across multiple VMs/containers.
        // This allows mounting XFS volumes that share the same UUID.
        if rootfs.fs_type == VM_ROOTFS_FILESYSTEM_XFS {
            // Add nouuid to existing options if not already present
            if !storage.options.iter().any(|opt| opt == "nouuid") {
                storage.options.push("nouuid".to_string());
            }
        }

        let mut device_id: String = "".to_owned();
        if let DeviceType::BlockModern(device_mod) = device_info {
            let device = device_mod.lock().await.clone();
            if block_driver == VIRTIO_PMEM {
                return Err(anyhow!(
                    "Complete support for block driver {} has not been implemented yet",
                    block_driver
                ));
            }

            storage.driver = device.config.driver_option.clone();
            storage.source = agent_storage_source_from_block_config(&device.config)?;
            device_id = device.device_id;
        }

        Ok(Self {
            guest_path: container_path.clone(),
            device_id,
            mount: oci::Mount::default(),
            storage: Some(storage),
        })
    }
}

#[async_trait]
impl Rootfs for BlockRootfs {
    async fn get_guest_rootfs_path(&self) -> Result<String> {
        Ok(self.guest_path.clone())
    }

    async fn get_rootfs_mount(&self) -> Result<Vec<oci::Mount>> {
        Ok(vec![self.mount.clone()])
    }

    async fn get_storage(&self) -> Option<Vec<Storage>> {
        self.storage.clone().map(|s| vec![s])
    }

    async fn get_device_id(&self) -> Result<Option<String>> {
        Ok(Some(self.device_id.clone()))
    }

    async fn cleanup(&self, device_manager: &RwLock<DeviceManager>) -> Result<()> {
        device_manager
            .write()
            .await
            .try_remove_device(&self.device_id)
            .await
    }
}

pub(crate) fn is_block_rootfs(m: &Mount) -> Option<(u64, Mount)> {
    if m.source.is_empty() {
        return None;
    }

    match stat::stat(m.source.as_str()) {
        Ok(fstat) => {
            if SFlag::from_bits_truncate(fstat.st_mode) == SFlag::S_IFBLK {
                let dev_id = fstat.st_rdev;
                let mut volume = m.clone();

                //clear the volume resource thus the block device will use the dev_id
                //to find the device's host path;
                volume.source = String::new();
                return Some((dev_id, volume));
            }

            if SFlag::from_bits_truncate(fstat.st_mode) == SFlag::S_IFREG
                && m.options.contains(&BLOCKFILE_ROOTFS_FLAG.to_string())
            {
                //use the block file's inode as the device id, which can make sure it's unique.
                let dev_id = fstat.st_ino;
                let options = m
                    .options
                    .clone()
                    .into_iter()
                    .filter(|o| !o.eq(BLOCKFILE_ROOTFS_FLAG))
                    .collect();

                //discard the blockfile rootfs's mount option "loop"
                let mut volume = m.clone();
                volume.options = options;

                return Some((dev_id, volume));
            }
        }
        Err(_) => return None,
    };
    None
}
