//! AArch64 architecture and platforms implementation

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

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

use crate::{
    arch::{aarch64::gic::Gic, Platform},
    device::{power::arm_psci::Psci, MACHINE_NAME},
    fs::{Initrd, INITRD_DATA},
    panic,
    util::call_init_array,
};

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

const BOOT_STACK_SIZE: usize = 4096 * 32;

/// Alias for terminal-level page tables
pub type L3 = mem::table::L3;

#[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() }
    }
}

static SHUTDOWN_FENCE: SpinFence = SpinFence::new();

impl Platform for AArch64 {
    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()
        }
    }

    unsafe fn power_off(&self) -> Result<!, Error> {
        Self::halt_aps()?;

        if let Some(psci) = self.psci.try_get() {
            log::info!("Powering off");
            psci.power_off()
        } else {
            log::warn!("No power off method, halting");
            ArchitectureImpl::halt();
        }
    }
}

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

    fn handle_ipi(ipi: Option<IpiMessage>) {
        let Some(ipi) = ipi else {
            log::warn!("Spurious IPI received, assuming panic");
            panic::panic_secondary();
        };

        match ipi {
            IpiMessage::Panic => panic::panic_secondary(),
            IpiMessage::Shutdown => Self::secondary_halt_handler(),
        }
    }

    fn secondary_halt_handler() {
        log::info!("CPU halted");
        SHUTDOWN_FENCE.signal();
        ArchitectureImpl::halt();
    }

    unsafe fn halt_aps() -> Result<(), Error> {
        let ap_count = CPU_COUNT.load(Ordering::Acquire) - 1;

        if ap_count > 0 {
            if let Some(gic) = GIC.try_get() {
                gic.send_ipi(IpiDeliveryTarget::OtherCpus, IpiMessage::Shutdown)?;
                SHUTDOWN_FENCE.wait_all(ap_count);
                Ok(())
            } else {
                log::warn!("No way to handle secondary CPUs");
                Err(Error::NotImplemented)
            }
        } else {
            Ok(())
        }
    }

    #[inline(never)]
    unsafe fn init_memory_management(
        &'static self,
        dtb_address: PhysicalAddress,
    ) -> Result<(), Error> {
        let dtb: PhysicalRef<[u8]> =
            PhysicalRef::map_slice(dtb_address, DeviceTree::MIN_HEADER_SIZE);
        let dtb_size = DeviceTree::read_totalsize(&dtb[..]).map_err(|_| Error::InvalidArgument)?;

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

        let dtb: PhysicalRef<[u8]> = PhysicalRef::map_slice(dtb_address, dtb_size);
        let dt = DeviceTree::from_raw(dtb.as_ptr().addr())?;

        // 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,
                },
            );
        }

        let dt = self.dt.init(dt);

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

        // 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(InitSequence::Early)?;
        }

        // 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 {
            atomic::compiler_fence(Ordering::SeqCst);

            debug::init();

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

            // Will register drivers
            call_init_array();

            // 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>"));
                },
                InitSequence::Early,
            );
            device_tree::driver::init_irqs(
                |_| (),
                |node, error| {
                    log::error!(
                        "{}: irq init error: {error:?}",
                        node.name().unwrap_or("<unknown>")
                    );
                },
                InitSequence::Early,
            );

            PciBusManager::probe_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(),
};
static GIC: OneTimeInit<Arc<Gic>> = OneTimeInit::new();
