// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! ARM64-specific topology definitions.

use super::ArchTopology;
use super::InvalidTopology;
use super::ProcessorTopology;
use super::TopologyBuilder;
use super::VpIndex;
use super::VpInfo;
use super::VpTopologyInfo;
use aarch64defs::MpidrEl1;

/// ARM64-specific topology information.
#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
#[derive(Debug, Copy, Clone)]
#[non_exhaustive]
pub struct Aarch64Topology {
    platform: Aarch64PlatformConfig,
}

impl ArchTopology for Aarch64Topology {
    type ArchVpInfo = Aarch64VpInfo;
    type BuilderState = Aarch64TopologyBuilderState;

    fn vp_topology(_topology: &ProcessorTopology<Self>, info: &Self::ArchVpInfo) -> VpTopologyInfo {
        VpTopologyInfo {
            socket: info.mpidr.aff2().into(),
            core: info.mpidr.aff1().into(),
            thread: info.mpidr.aff0().into(),
        }
    }
}

/// Aarch64-specific [`TopologyBuilder`] state.
pub struct Aarch64TopologyBuilderState {
    platform: Aarch64PlatformConfig,
}

/// ARM64 platform interrupt and GIC configuration.
///
/// Groups GIC base addresses, MSI frame info, and platform interrupt
/// assignments (PMU, virtual timer) into a single struct so that the
/// topology builder takes one value instead of several positional `u32`s.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
pub struct Aarch64PlatformConfig {
    /// GIC distributor base address.
    #[cfg_attr(feature = "inspect", inspect(hex))]
    pub gic_distributor_base: u64,
    /// GIC redistributors base address.
    #[cfg_attr(feature = "inspect", inspect(hex))]
    pub gic_redistributors_base: u64,
    /// GIC v2m MSI frame, if MSIs via v2m are supported.
    pub gic_v2m: Option<GicV2mInfo>,
    /// Performance Monitor Unit GSIV (GIC INTID). `None` if not available.
    pub pmu_gsiv: Option<u32>,
    /// Virtual timer PPI (GIC INTID, e.g. 20 for PPI 4).
    pub virt_timer_ppi: u32,
}

/// GIC v2m MSI frame parameters.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
pub struct GicV2mInfo {
    /// Physical base address of the v2m MSI frame.
    #[cfg_attr(feature = "inspect", inspect(hex))]
    pub frame_base: u64,
    /// First GIC interrupt ID in the SPI range owned by this frame.
    pub spi_base: u32,
    /// Number of SPIs owned by this frame.
    pub spi_count: u32,
}

/// ARM64 specific VP info.
#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
#[derive(Debug, Copy, Clone)]
pub struct Aarch64VpInfo {
    /// The base info.
    #[cfg_attr(feature = "inspect", inspect(flatten))]
    pub base: VpInfo,
    /// The MPIDR_EL1 value of the processor.
    #[cfg_attr(feature = "inspect", inspect(hex, with = "|&x| u64::from(x)"))]
    pub mpidr: MpidrEl1,
    /// GIC Redistributor Address
    #[cfg_attr(feature = "inspect", inspect(hex))]
    pub gicr: u64,
    /// Performance Interrupt GSIV (PMU)
    #[cfg_attr(feature = "inspect", inspect(hex))]
    pub pmu_gsiv: Option<u32>,
}

impl AsRef<VpInfo> for Aarch64VpInfo {
    fn as_ref(&self) -> &VpInfo {
        &self.base
    }
}

impl TopologyBuilder<Aarch64Topology> {
    /// Returns a builder for creating an aarch64 processor topology.
    pub fn new_aarch64(platform: Aarch64PlatformConfig) -> Self {
        Self {
            vps_per_socket: 1,
            smt_enabled: false,
            arch: Aarch64TopologyBuilderState { platform },
        }
    }

    /// Builds a processor topology with `proc_count` processors.
    pub fn build(
        &self,
        proc_count: u32,
    ) -> Result<ProcessorTopology<Aarch64Topology>, InvalidTopology> {
        if proc_count >= 256 {
            return Err(InvalidTopology::TooManyVps {
                requested: proc_count,
                max: u8::MAX.into(),
            });
        }
        if !(16..32).contains(&self.arch.platform.virt_timer_ppi) {
            return Err(InvalidTopology::InvalidPpiIntid(
                self.arch.platform.virt_timer_ppi,
            ));
        }
        if let Some(gsiv) = self.arch.platform.pmu_gsiv {
            if !(16..32).contains(&gsiv) {
                return Err(InvalidTopology::InvalidPpiIntid(gsiv));
            }
        }
        let mpidrs = (0..proc_count).map(|vp_index| {
            // TODO: construct mpidr appropriately for the specified
            // topology.
            let uni_proc = proc_count == 1;
            let mut aff = (0..4).map(|i| (vp_index >> (8 * i)) as u8);
            MpidrEl1::new()
                .with_res1_31(true)
                .with_u(uni_proc)
                .with_aff0(aff.next().unwrap())
                .with_aff1(aff.next().unwrap())
                .with_aff2(aff.next().unwrap())
                .with_aff3(aff.next().unwrap())
        });
        self.build_with_vp_info(mpidrs.enumerate().map(|(id, mpidr)| Aarch64VpInfo {
            base: VpInfo {
                vp_index: VpIndex::new(id as u32),
                vnode: 0,
            },
            mpidr,
            gicr: self.arch.platform.gic_redistributors_base
                + id as u64 * aarch64defs::GIC_REDISTRIBUTOR_SIZE,
            pmu_gsiv: self.arch.platform.pmu_gsiv,
        }))
    }

    /// Builds a processor topology with processors with the specified information.
    pub fn build_with_vp_info(
        &self,
        vps: impl IntoIterator<Item = Aarch64VpInfo>,
    ) -> Result<ProcessorTopology<Aarch64Topology>, InvalidTopology> {
        let vps = Vec::from_iter(vps);
        let mut smt_enabled = false;
        for (i, vp) in vps.iter().enumerate() {
            if i != vp.base.vp_index.index() as usize {
                return Err(InvalidTopology::InvalidVpIndices);
            }

            if vp.mpidr.mt() {
                smt_enabled = true;
            }
        }

        Ok(ProcessorTopology {
            vps,
            smt_enabled,
            vps_per_socket: self.vps_per_socket,
            arch: Aarch64Topology {
                platform: self.arch.platform,
            },
        })
    }
}

impl ProcessorTopology<Aarch64Topology> {
    /// Returns the GIC distributor base
    pub fn gic_distributor_base(&self) -> u64 {
        self.arch.platform.gic_distributor_base
    }

    /// Returns the GIC redistributors base
    pub fn gic_redistributors_base(&self) -> u64 {
        self.arch.platform.gic_redistributors_base
    }

    /// Returns the PMU GSIV
    pub fn pmu_gsiv(&self) -> Option<u32> {
        self.arch.platform.pmu_gsiv
    }

    /// Returns the GIC v2m MSI frame info, if present.
    pub fn gic_v2m(&self) -> Option<GicV2mInfo> {
        self.arch.platform.gic_v2m
    }

    /// Returns the virtual timer PPI (GIC INTID).
    pub fn virt_timer_ppi(&self) -> u32 {
        self.arch.platform.virt_timer_ppi
    }
}
