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

use std::{collections::HashMap, convert::TryFrom};

use anyhow::{ensure, Context, Result};
use kata_types::{
    annotations::Annotation, config::TomlConfig, container::ContainerType,
    cpu::LinuxContainerCpuResources, k8s::container_type,
};
use oci_spec::runtime as oci;

// initial resource that InitialSizeManager needs, this is the spec for the
// sandbox/container's workload
#[derive(Clone, Copy, Debug)]
struct InitialSize {
    vcpu: f32,
    mem_mb: u32,
    orig_toml_default_mem: u32,
}

const MIB: i64 = 1024 * 1024;

// generate initial resource(vcpu and memory in MiB) from annotations
impl TryFrom<&HashMap<String, String>> for InitialSize {
    type Error = anyhow::Error;
    fn try_from(an: &HashMap<String, String>) -> Result<Self> {
        let mut vcpu: f32 = 0.0;

        let annotation = Annotation::new(an.clone());
        let (period, quota, memory) =
            get_sizing_info(annotation).context("failed to get sizing info")?;
        let mut cpu = oci::LinuxCpu::default();
        cpu.set_period(Some(period));
        cpu.set_quota(Some(quota));

        // although it may not be actually a linux container, we are only using the calculation inside
        // LinuxContainerCpuResources::try_from to generate our vcpu number
        if let Ok(cpu_resource) = LinuxContainerCpuResources::try_from(&cpu) {
            vcpu = get_nr_vcpu(&cpu_resource);
        }
        let mem_mb = convert_memory_to_mb(memory);

        Ok(Self {
            vcpu,
            mem_mb,
            orig_toml_default_mem: 0,
        })
    }
}

// generate initial resource(vcpu and memory in MiB) from spec's information
impl TryFrom<&oci::Spec> for InitialSize {
    type Error = anyhow::Error;
    fn try_from(spec: &oci::Spec) -> Result<Self> {
        let mut vcpu: f32 = 0.0;
        let mut mem_mb: u32 = 0;
        match container_type(spec) {
            // podsandbox, from annotation
            ContainerType::PodSandbox => {
                let spec_annos = spec.annotations().clone().unwrap_or_default();
                return InitialSize::try_from(&spec_annos);
            }
            // single container, from container spec
            _ => {
                if let Some(resource) = spec
                    .linux()
                    .as_ref()
                    .and_then(|linux| linux.resources().as_ref())
                {
                    // cpu resource
                    if let Some(Ok(cpu_resource)) = resource
                        .cpu()
                        .as_ref()
                        .map(LinuxContainerCpuResources::try_from)
                    {
                        vcpu = get_nr_vcpu(&cpu_resource);
                    }

                    // memory resource
                    mem_mb = resource
                        .memory()
                        .as_ref()
                        .and_then(|mem| mem.limit())
                        .map(convert_memory_to_mb)
                        .unwrap_or(0);
                }
            }
        }
        info!(
            sl!(),
            "(from PodSandbox's annotation / SingleContainer's spec) initial size: vcpu={}, mem_mb={}", vcpu, mem_mb
        );
        Ok(Self {
            vcpu,
            mem_mb,
            orig_toml_default_mem: 0,
        })
    }
}

// InitialSizeManager is responsible for initial vcpu/mem management
//
// inital vcpu/mem management sizing information is optionally provided, either by
// upper layer runtime (containerd / crio) or by the container spec itself (when it
// is a standalone single container such as the one started with *docker run*)
//
// the sizing information uses three values, cpu quota, cpu period and memory limit,
// and with above values it calculates the # vcpus and memory for the workload
//
// if the workload # of vcpus and memory is invalid for vmms, we still use default
// value in toml_config
#[derive(Clone, Copy, Debug)]
pub struct InitialSizeManager {
    resource: InitialSize,
}

impl InitialSizeManager {
    pub fn new(spec: &oci::Spec) -> Result<Self> {
        Ok(Self {
            resource: InitialSize::try_from(spec).context("failed to construct static resource")?,
        })
    }

    pub fn new_from(annotation: &HashMap<String, String>) -> Result<Self> {
        Ok(Self {
            resource: InitialSize::try_from(annotation)
                .context("failed to construct static resource")?,
        })
    }

    // Merge sizing values from sandbox annotations when the current source
    // (typically the OCI spec) does not carry CRI sandbox sizing keys.
    pub fn supplement_from_annotations(
        &mut self,
        annotation: &HashMap<String, String>,
    ) -> Result<()> {
        let from_annotation =
            InitialSize::try_from(annotation).context("failed to construct static resource")?;

        if self.resource.vcpu == 0.0 {
            self.resource.vcpu = from_annotation.vcpu;
        }
        if self.resource.mem_mb == 0 {
            self.resource.mem_mb = from_annotation.mem_mb;
        }

        Ok(())
    }

    pub fn setup_config(&mut self, config: &mut TomlConfig) -> Result<()> {
        // update this data to the hypervisor config for later use by hypervisor
        let hypervisor_name = &config.runtime.hypervisor_name;
        let hv = config
            .hypervisor
            .get_mut(hypervisor_name)
            .context("failed to get hypervisor config")?;

        self.resource.orig_toml_default_mem = hv.memory_info.default_memory;

        // Non-static mode keeps configured defaults unchanged.
        if !config.runtime.static_sandbox_resource_mgmt {
            validate_non_zero_sandbox_memory(hypervisor_name, hv.memory_info.default_memory)?;
            return Ok(());
        }

        if self.resource.vcpu > 0.0 || self.resource.mem_mb > 0 {
            if self.resource.vcpu > 0.0 {
                info!(sl!(), "resource with vcpu {}", self.resource.vcpu);
            }
            if self.resource.mem_mb > 0 {
                info!(sl!(), "resource with memory {}", self.resource.mem_mb);
            }

            hv.cpu_info.default_vcpus = (hv.cpu_info.overhead_vcpus + self.resource.vcpu).max(1.0);

            hv.memory_info.default_memory = hv.memory_info.overhead_memory + self.resource.mem_mb;
            hv.memory_info.default_maxmemory = hv
                .memory_info
                .default_maxmemory
                .max(hv.memory_info.default_memory);
        }

        hv.cpu_info.default_maxvcpus = hv.cpu_info.default_vcpus.ceil() as u32;

        validate_non_zero_sandbox_memory(hypervisor_name, hv.memory_info.default_memory)?;
        Ok(())
    }

    pub fn get_orig_toml_default_mem(&self) -> u32 {
        self.resource.orig_toml_default_mem
    }
}

fn validate_non_zero_sandbox_memory(hypervisor_name: &str, memory_mib: u32) -> Result<()> {
    ensure!(
        memory_mib > 0,
        "computed sandbox memory is 0 MiB for hypervisor '{}'; set a non-zero memory limit or configure non-zero default_memory/overhead_memory",
        hypervisor_name
    );
    Ok(())
}

fn get_nr_vcpu(resource: &LinuxContainerCpuResources) -> f32 {
    if let Some(v) = resource.get_vcpus() {
        v as f32
    } else {
        0.0
    }
}

fn convert_memory_to_mb(memory_in_byte: i64) -> u32 {
    if memory_in_byte < 0 {
        return 0;
    }
    let mem_size = (memory_in_byte / MIB) as u32;
    // memory size must be 2MB aligned for hugepage support
    if !mem_size.is_multiple_of(2) {
        return mem_size + 1;
    }

    mem_size
}

// from the upper layer runtime's annotation (e.g. crio, k8s), get the *cpu quota,
// cpu period and memory limit* for a sandbox/container
fn get_sizing_info(annotation: Annotation) -> Result<(u64, i64, i64)> {
    // since we are *adding* our result to the config, a value of 0 will cause no change
    // and if the annotation is not assigned (but static resource management is), we will
    // log a *warning* to fill that with zero value
    let period = annotation.get_sandbox_cpu_period();
    let quota = annotation.get_sandbox_cpu_quota();
    let memory = annotation.get_sandbox_mem();
    Ok((period, quota, memory))
}

#[cfg(test)]
mod tests {
    use super::*;
    use kata_types::annotations::cri_containerd;
    use oci_spec::runtime::{LinuxBuilder, LinuxMemory, LinuxMemoryBuilder, LinuxResourcesBuilder};
    use rstest::rstest;
    use std::collections::HashMap;
    #[derive(Clone)]
    struct InputData {
        period: Option<u64>,
        quota: Option<i64>,
        memory: Option<i64>,
    }

    #[derive(Clone)]
    struct TestData<'a> {
        desc: &'a str,
        input: InputData,
        result: InitialSize,
    }

    fn get_test_data() -> Vec<TestData<'static>> {
        [
            TestData {
                desc: "no resource limit",
                input: InputData {
                    period: None,
                    quota: None,
                    memory: None,
                },
                result: InitialSize {
                    vcpu: 0.0,
                    mem_mb: 0,
                    orig_toml_default_mem: 0,
                },
            },
            TestData {
                desc: "normal resource limit",
                // data below should result in 2200 mCPU(round up to 3 vcpus) and 512 MiB of memory
                input: InputData {
                    period: Some(100_000),
                    quota: Some(220_000),
                    memory: Some(512 * MIB),
                },
                result: InitialSize {
                    vcpu: 3.0,
                    mem_mb: 512,
                    orig_toml_default_mem: 0,
                },
            },
            TestData {
                desc: "Odd memory in resource limits",
                input: InputData {
                    period: None,
                    quota: None,
                    memory: Some(513 * MIB),
                },
                result: InitialSize {
                    vcpu: 0.0,
                    mem_mb: 514,
                    orig_toml_default_mem: 0,
                },
            },
        ]
        .to_vec()
    }

    #[test]
    fn test_initial_size_sandbox() {
        let tests = get_test_data();

        // run tests
        for (i, d) in tests.iter().enumerate() {
            let mut spec = oci::Spec::default();
            spec.set_annotations(Some(HashMap::from([
                (
                    cri_containerd::CONTAINER_TYPE_LABEL_KEY.to_string(),
                    cri_containerd::SANDBOX.to_string(),
                ),
                (
                    cri_containerd::SANDBOX_CPU_PERIOD_KEY.to_string(),
                    d.input.period.map_or(String::new(), |v| format!("{}", v)),
                ), // CPU period
                (
                    cri_containerd::SANDBOX_CPU_QUOTA_KEY.to_string(),
                    d.input.quota.map_or(String::new(), |v| format!("{}", v)),
                ), // CPU quota
                (
                    cri_containerd::SANDBOX_MEM_KEY.to_string(),
                    d.input.memory.map_or(String::new(), |v| format!("{}", v)),
                ), // memory in bytes
            ])));

            let initial_size = InitialSize::try_from(&spec);
            assert!(
                initial_size.is_ok(),
                "test[{}]: {:?} should be ok",
                i,
                d.desc
            );

            let initial_size = initial_size.unwrap();
            assert_eq!(
                initial_size.vcpu.ceil(),
                d.result.vcpu,
                "test[{}]: {:?} vcpu should be {}",
                i,
                d.desc,
                d.result.vcpu,
            );
            assert_eq!(
                initial_size.mem_mb, d.result.mem_mb,
                "test[{}]: {:?} memory should be {}",
                i, d.desc, d.result.mem_mb,
            );
        }
    }

    #[test]
    fn test_initial_size_container() {
        let tests = get_test_data();

        // run tests
        for (i, d) in tests.iter().enumerate() {
            let mut spec = oci::Spec::default();
            spec.set_annotations(Some(HashMap::from([(
                cri_containerd::CONTAINER_TYPE_LABEL_KEY.to_string(),
                cri_containerd::CONTAINER.to_string(),
            )])));

            let mut linux_cpu = oci::LinuxCpu::default();
            linux_cpu.set_period(d.input.period);
            linux_cpu.set_quota(d.input.quota);
            let mut linux_mem: LinuxMemory = LinuxMemory::default();
            if let Some(limit) = d.input.memory {
                linux_mem = LinuxMemoryBuilder::default().limit(limit).build().unwrap();
            };

            let linux = LinuxBuilder::default()
                .resources(
                    LinuxResourcesBuilder::default()
                        .cpu(linux_cpu)
                        .memory(linux_mem)
                        .build()
                        .unwrap(),
                )
                .build()
                .unwrap();
            spec.set_linux(Some(linux));

            let initial_size = InitialSize::try_from(&spec);
            assert!(
                initial_size.is_ok(),
                "test[{}]: {:?} should be ok",
                i,
                d.desc
            );

            let initial_size = initial_size.unwrap();
            assert_eq!(
                initial_size.vcpu.ceil(),
                d.result.vcpu,
                "test[{}]: {:?} vcpu should be {}",
                i,
                d.desc,
                d.result.vcpu,
            );
            assert_eq!(
                initial_size.mem_mb, d.result.mem_mb,
                "test[{}]: {:?} memory should be {}",
                i, d.desc, d.result.mem_mb,
            );
        }
    }

    fn make_config(
        default_vcpus: f32,
        overhead_vcpus: f32,
        default_maxvcpus: u32,
        default_memory: u32,
        overhead_memory: u32,
        default_maxmemory: u32,
        static_sandbox_resource_mgmt: bool,
    ) -> TomlConfig {
        use kata_types::config::Hypervisor;

        let mut config = TomlConfig::default();
        config
            .hypervisor
            .insert("qemu".to_owned(), Hypervisor::default());
        config.hypervisor.entry("qemu".to_owned()).and_modify(|hv| {
            hv.cpu_info.default_vcpus = default_vcpus;
            hv.cpu_info.overhead_vcpus = overhead_vcpus;
            hv.cpu_info.default_maxvcpus = default_maxvcpus;
            hv.memory_info.default_memory = default_memory;
            hv.memory_info.overhead_memory = overhead_memory;
            hv.memory_info.default_maxmemory = default_maxmemory;
        });
        config.runtime.hypervisor_name = "qemu".to_owned();
        config.runtime.static_sandbox_resource_mgmt = static_sandbox_resource_mgmt;
        config
    }

    #[test]
    fn test_setup_config_static_applies_vcpu_and_memory() {
        let mut config = make_config(1.0, 0.5, 4, 256, 128, 4096, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 1.2,
                mem_mb: 512,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();
        assert_eq!(hv.cpu_info.default_vcpus, 1.7);
        assert_eq!(hv.memory_info.default_memory, 640);
    }

    #[test]
    fn test_setup_config_non_static_does_not_apply() {
        let mut config = make_config(1.0, 0.5, 4, 256, 128, 4096, false);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 1.2,
                mem_mb: 512,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();
        assert_eq!(hv.cpu_info.default_vcpus, 1.0);
        assert_eq!(hv.memory_info.default_memory, 256);
    }

    #[test]
    fn test_setup_config_clamps_maxvcpus() {
        let mut config = make_config(1.0, 1.0, 2, 256, 128, 4096, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 2.5,
                mem_mb: 0,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();
        assert_eq!(hv.cpu_info.default_vcpus, 3.5);
        assert_eq!(hv.cpu_info.default_maxvcpus, 4);
    }

    #[test]
    fn test_setup_config_static_reduces_maxvcpus_to_static_total() {
        let mut config = make_config(1.0, 0.5, 8, 256, 128, 4096, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 1.2,
                mem_mb: 0,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();
        assert_eq!(hv.cpu_info.default_vcpus, 1.7);
        assert_eq!(hv.cpu_info.default_maxvcpus, 2);
    }

    #[test]
    fn test_setup_config_clamps_maxmemory() {
        let mut config = make_config(1.0, 0.5, 4, 256, 128, 300, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 0.0,
                mem_mb: 512,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();
        assert_eq!(hv.memory_info.default_memory, 640);
        assert_eq!(hv.memory_info.default_maxmemory, 640);
    }

    #[test]
    fn test_setup_config_preserves_orig_toml_default_mem() {
        let mut config = make_config(1.0, 0.5, 4, 256, 128, 4096, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 0.0,
                mem_mb: 128,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        assert_eq!(mgr.get_orig_toml_default_mem(), 256);
    }

    #[test]
    fn test_supplement_from_annotations_fills_missing_spec_sizing() {
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 0.0,
                mem_mb: 0,
                orig_toml_default_mem: 0,
            },
        };

        let ann = HashMap::from([
            (
                cri_containerd::SANDBOX_CPU_PERIOD_KEY.to_string(),
                "100000".to_string(),
            ),
            (
                cri_containerd::SANDBOX_CPU_QUOTA_KEY.to_string(),
                "120000".to_string(),
            ),
            (
                cri_containerd::SANDBOX_MEM_KEY.to_string(),
                (256 * MIB).to_string(),
            ),
        ]);

        mgr.supplement_from_annotations(&ann).unwrap();

        const VCPU_TOLERANCE: f32 = 0.0001;
        assert!((mgr.resource.vcpu - 1.2).abs() < VCPU_TOLERANCE);
        assert_eq!(mgr.resource.mem_mb, 256);
    }

    #[test]
    fn test_setup_config_static_without_limits_uses_toml_defaults() {
        let mut config = make_config(2.0, 0.5, 8, 512, 128, 4096, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 0.0,
                mem_mb: 0,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();
        assert_eq!(hv.cpu_info.default_vcpus, 2.0);
        assert_eq!(hv.memory_info.default_memory, 512);
    }

    #[test]
    fn test_setup_config_static_errors_on_zero_memory() {
        let mut config = make_config(1.0, 0.5, 8, 1024, 0, 4096, true);
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: 1.0,
                mem_mb: 0,
                orig_toml_default_mem: 0,
            },
        };

        let err = mgr.setup_config(&mut config).unwrap_err().to_string();
        assert!(err.contains("computed sandbox memory is 0 MiB"));
        assert!(err.contains("default_memory/overhead_memory"));
    }

    #[rstest]
    #[case::both_limits(3.0, 0.75, 1024, 256, 1.25, 1024, 2.0, 1280)]
    #[case::cpu_only_limit(3.0, 0.5, 1024, 128, 1.5, 0, 2.0, 128)]
    #[case::memory_only_limit(3.0, 0.5, 1024, 128, 0.0, 512, 1.0, 640)]
    #[case::both_limits_zero_overhead(3.0, 0.0, 1024, 0, 1.25, 1024, 1.25, 1024)]
    #[case::memory_only_zero_overhead(3.0, 0.0, 1024, 0, 0.0, 512, 1.0, 512)]
    fn test_setup_config_static_requested_vs_defaults(
        #[case] default_vcpus: f32,
        #[case] overhead_vcpus: f32,
        #[case] default_memory: u32,
        #[case] overhead_memory: u32,
        #[case] requested_vcpus: f32,
        #[case] requested_mem_mb: u32,
        #[case] expected_default_vcpus: f32,
        #[case] expected_default_memory: u32,
    ) {
        let mut config = make_config(
            default_vcpus,
            overhead_vcpus,
            8,
            default_memory,
            overhead_memory,
            4096,
            true,
        );
        let mut mgr = InitialSizeManager {
            resource: InitialSize {
                vcpu: requested_vcpus,
                mem_mb: requested_mem_mb,
                orig_toml_default_mem: 0,
            },
        };

        mgr.setup_config(&mut config).unwrap();
        let hv = config.hypervisor.get("qemu").unwrap();

        assert_eq!(hv.cpu_info.default_vcpus, expected_default_vcpus);
        assert_eq!(hv.memory_info.default_memory, expected_default_memory);
    }
}
