//! AArch64 architecture and platforms implementation

use core::sync::atomic::{self, Ordering};

use aarch64_cpu::registers::{CNTP_CTL_EL0, CNTP_TVAL_EL0};
use abi::error::Error;
use alloc::sync::Arc;
use device_api::{
    interrupt::{Irq, LocalInterruptController},
    ResetDevice,
};
use device_tree::{driver::unflatten_device_tree, DeviceTree, DeviceTreeNodeExt};
use kernel_arch_aarch64::{
    mem::{
        self,
        table::{L1, L3},
        EarlyMapping, MEMORY_LIMIT, RAM_MAPPING_L1_COUNT,
    },
    ArchitectureImpl, PerCpuData,
};
use libk::{arch::Cpu, config, debug, device::external_interrupt_controller};
use libk_mm::{
    address::PhysicalAddress,
    phys::{self, reserved::reserve_region, PhysicalMemoryRegion},
    pointer::PhysicalRef,
    table::EntryLevelExt,
};
use libk_util::OneTimeInit;
use tock_registers::interfaces::Writeable;
use ygg_driver_pci::PciBusManager;

use crate::{
    // device::power::arm_psci::Psci,
    device::{power::arm_psci::Psci, MACHINE_NAME},
    fs::{Initrd, INITRD_DATA},
    util::call_init_array,
};

use self::gic::Gic;

use super::Platform;

pub mod boot;
pub mod exception;
pub mod gic;
pub mod smp;
pub mod timer;

const BOOT_STACK_SIZE: usize = 4096 * 32;

#[derive(Clone, Copy)]
#[repr(C, align(0x20))]
struct BootStack<const SIZE: usize> {
    data: [u8; SIZE],
}

/// AArch64 architecture implementation
pub struct AArch64 {
    dt: OneTimeInit<DeviceTree<'static>>,

    /// Optional instance of PSCI on this platform
    pub psci: OneTimeInit<Arc<Psci>>,
    reset: OneTimeInit<Arc<dyn ResetDevice>>,
    initrd: OneTimeInit<PhysicalRef<'static, [u8]>>,
    machine_compatible: OneTimeInit<&'static str>,
}

impl<const SIZE: usize> BootStack<SIZE> {
    pub const fn zeroed() -> Self {
        Self { data: [0; SIZE] }
    }

    pub fn top_addr(&self) -> usize {
        unsafe { self.data.as_ptr().add(SIZE).addr() }
    }
}

impl Platform for AArch64 {
    const KERNEL_VIRT_OFFSET: usize = 0xFFFFFF8000000000;

    type L3 = L3;

    unsafe fn start_application_processors(&self) {
        if let Some(compatible) = self.machine_compatible.try_get() {
            if *compatible == "raspberrypi,4-model-b" {
                log::warn!("raspi4b: cache workaround disable SMP");
                return;
            }
        }

        let dt = self.dt.get();
        if let Err(error) = smp::start_ap_cores(dt) {
            log::error!("Could not initialize AP CPUs: {:?}", error);
        }
    }

    fn register_reset_device(&self, reset: Arc<dyn ResetDevice>) -> Result<(), Error> {
        self.reset.init(reset);
        Ok(())
    }

    unsafe fn reset(&self) -> ! {
        if let Some(reset) = self.reset.try_get() {
            reset.reset()
        } else {
            let psci = self.psci.get();
            psci.reset()
        }
    }
}

static GIC: OneTimeInit<Arc<Gic>> = OneTimeInit::new();

impl AArch64 {
    fn set_gic(gic: Arc<Gic>) {
        GIC.init(gic);
    }

    fn map_physical_memory<I: Iterator<Item = PhysicalMemoryRegion> + Clone>(
        _it: I,
        _memory_start: PhysicalAddress,
        memory_end: PhysicalAddress,
    ) -> Result<(), Error> {
        let end_l1i = memory_end.page_align_up::<L1>().page_index::<L1>();
        if end_l1i > RAM_MAPPING_L1_COUNT {
            panic!("TODO: partial physical memory mapping");
        }

        // Map 1GiB chunks
        for index in 0..end_l1i {
            unsafe {
                kernel_arch_aarch64::mem::map_ram_l1(index);
            }
        }

        MEMORY_LIMIT.store(memory_end.into_usize(), Ordering::Release);

        Ok(())
    }

    unsafe fn init_memory_management(&'static self, dtb: PhysicalAddress) -> Result<(), Error> {
        // Initialize the runtime mappings
        kernel_arch_aarch64::mem::init_fixed_tables();

        // Extract the size of the device tree
        let dtb_size = {
            let dtb_header = EarlyMapping::<u8>::map_slice(dtb, DeviceTree::MIN_HEADER_SIZE)?;
            DeviceTree::read_totalsize(dtb_header.as_ref()).map_err(|_| Error::InvalidArgument)?
        };

        reserve_region(
            "dtb",
            PhysicalMemoryRegion {
                base: dtb,
                size: (dtb_size + 0xFFF) & !0xFFF,
            },
        );

        let dtb_slice = EarlyMapping::<u8>::map_slice(dtb, dtb_size)?;
        let dt = DeviceTree::from_raw(dtb_slice.as_ptr() as usize)?;

        // Setup initrd from the dt
        let initrd = dt.chosen_initrd();

        if let Some((start, end)) = initrd {
            let aligned_start = start.page_align_down::<L3>();
            let aligned_end = end.page_align_up::<L3>();

            let size = aligned_end - aligned_start;
            reserve_region(
                "initrd",
                PhysicalMemoryRegion {
                    base: aligned_start,
                    size,
                },
            );
        }

        // Initialize the physical memory
        phys::init_from_iter(dt.memory_regions(), Self::map_physical_memory)?;

        // EarlyMapping for DTB no longer needed, it lives in physical memory and can be obtained
        // through PhysicalRef
        let dtb_slice: PhysicalRef<'static, [u8]> = PhysicalRef::map_slice(dtb, dtb_size);
        let dt = DeviceTree::from_raw(dtb_slice.as_ptr().addr())?;

        self.dt.init(dt);

        // Setup initrd
        if let Some((initrd_start, initrd_end)) = initrd {
            let aligned_start = initrd_start.page_align_down::<L3>();
            let aligned_end = initrd_end.page_align_up::<L3>();
            let len = initrd_end - initrd_start;

            let data = unsafe { PhysicalRef::map_slice(initrd_start, len) };
            self.initrd.init(data);

            INITRD_DATA.init(Initrd {
                phys_page_start: aligned_start,
                phys_page_len: aligned_end - aligned_start,
                data: self.initrd.get().as_ref(),
            });
        }

        Ok(())
    }

    #[inline(never)]
    unsafe fn setup_chosen_stdout(dt: &'static DeviceTree) -> Result<(), Error> {
        // Get /chosen.stdout-path to get early debug printing
        // TODO honor defined configuration value
        let stdout = dt.chosen_stdout();
        let stdout_path = stdout.map(|(p, _)| p);
        let node = stdout_path.and_then(device_tree::driver::find_node);

        if let Some(node) = node {
            node.force_init()?;
        }

        // No stdout
        Ok(())
    }

    fn machine_name(dt: &'static DeviceTree) -> (Option<&'static str>, Option<&'static str>) {
        (
            dt.root().prop_string("compatible"),
            dt.root().prop_string("model"),
        )
    }

    fn apply_machine_workarounds(compatible: &str) {
        #[allow(clippy::single_match)]
        match compatible {
            "raspberrypi,4-model-b" => {
                // XXX Workaround for Raspberry Pi 4B:
                //  Instruction cache totally freaks out at EL0 and I don't know why
                unsafe { mem::disable_icache() };
            }
            _ => (),
        }
    }

    unsafe fn init_platform(&'static self, is_bsp: bool) -> Result<(), Error> {
        let per_cpu = PerCpuData {
            gic: OneTimeInit::new(),
        };
        Cpu::init_local(None, per_cpu);

        if is_bsp {
            // Will register drivers
            call_init_array();

            atomic::compiler_fence(Ordering::SeqCst);

            debug::init();

            let dt = self.dt.get();
            let bootargs = dt.chosen_bootargs().unwrap_or("");
            config::parse_boot_arguments(bootargs);

            // Create device tree sysfs nodes
            device_tree::util::create_sysfs_nodes(dt);

            let (machine_compatible, machine_name) = Self::machine_name(dt);

            if let Some(compatible) = machine_compatible {
                self.machine_compatible.init(compatible);
                Self::apply_machine_workarounds(compatible);
            }

            unflatten_device_tree(dt);

            Self::setup_chosen_stdout(dt).ok();

            if let Some(machine) = machine_name {
                log::info!("Running on {machine:?}");
                MACHINE_NAME.init(machine.into());
            }
            log::info!("Boot arguments: {bootargs:?}");

            log::info!("Initializing aarch64 platform");

            device_tree::driver::lazy_init(
                |_| (),
                |node, error| {
                    log::error!("{}: {error:?}", node.name().unwrap_or("<unknown>"));
                },
            );
            device_tree::driver::init_irqs(
                |_| (),
                |node, error| {
                    log::error!(
                        "{}: irq init error: {error:?}",
                        node.name().unwrap_or("<unknown>")
                    );
                },
            );

            PciBusManager::setup_bus_devices()?;
        } else {
            // BSP already initialized everything needed
            // Setup timer and local interrupt controller
            if let Some(gic) = GIC.try_get() {
                unsafe {
                    gic.init_ap().unwrap();
                }
            }

            // TODO device-tree initialization for this
            CNTP_CTL_EL0.write(CNTP_CTL_EL0::ENABLE::SET + CNTP_CTL_EL0::IMASK::CLEAR);
            CNTP_TVAL_EL0.set(10000000);
            external_interrupt_controller()
                .unwrap()
                .enable_irq(Irq::Private(14))
                .unwrap();
        }

        if let Some(gic) = GIC.try_get() {
            let cpu_data = ArchitectureImpl::local_cpu_data().unwrap();
            cpu_data.gic.init(gic.clone());
        }

        Ok(())
    }
}

/// AArch64 implementation value
pub static PLATFORM: AArch64 = AArch64 {
    dt: OneTimeInit::new(),
    initrd: OneTimeInit::new(),

    psci: OneTimeInit::new(),
    reset: OneTimeInit::new(),
    machine_compatible: OneTimeInit::new(),
};