//! 64-bit RISC-V platform implementation
use core::sync::atomic::{self, Ordering};

use abi::error::Error;
use device_api::clock::Hertz;
use device_tree::{
    driver::{unflatten_device_tree, InitSequence},
    DeviceTree, DeviceTreeNodeExt,
};
use kernel_arch::{Architecture, ArchitectureImpl};
use kernel_arch_riscv64::{
    mem,
    registers::{SIE, SSTATUS},
    sbi, PerCpuData,
};
use libk::{arch::Cpu, config};
use libk_mm::{
    address::PhysicalAddress,
    phys::{self, reserved::reserve_region, PhysicalMemoryRegion},
    pointer::PhysicalRef,
    table::EntryLevelExt,
};
use libk_util::OneTimeInit;
use tock_registers::interfaces::ReadWriteable;
use ygg_driver_pci::PciBusManager;

use crate::{
    arch::Platform,
    device::MACHINE_NAME,
    fs::{Initrd, INITRD_DATA},
    util::call_init_array,
};

pub mod boot;
pub mod debug;
pub mod exception;
pub mod smp;
pub mod timer;

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

/// 64-bit RISC-V platform struct
pub struct Riscv64 {
    dt: OneTimeInit<DeviceTree<'static>>,
    initrd: OneTimeInit<PhysicalRef<'static, [u8]>>,
}

impl Platform for Riscv64 {
    unsafe fn reset(&self) -> ! {
        ArchitectureImpl::halt();
    }

    unsafe fn power_off(&self) -> Result<!, Error> {
        sbi::sbi_system_shutdown()
    }

    unsafe fn start_application_processors(&self) {
        // TODO asymmetric systems with different hart types are not yet supported.
        // e.g., in JH7110 there're two different types of cores
        let dt = self.dt.get();
        if let Err(error) = smp::start_secondary_harts(dt) {
            log::error!("Couldn't start secondary harts: {error:?}");
        }
    }
}

/// 64-bit RISC-V platform
pub static PLATFORM: Riscv64 = Riscv64 {
    dt: OneTimeInit::new(),
    initrd: OneTimeInit::new(),
};

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

        //         // Unmap the lower half
        //         mem::setup_fixed_tables();

        log::info!("DTB: {:#x?}", dtb_address..dtb_address.add(dtb_size));

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

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

        // Reserve memory regions specified in the DTB
        log::info!("Reserved memory:");
        for region in dt.reserved_regions() {
            log::info!("* {:#x}..{:#x}", region.base, region.base.add(region.size));
            reserve_region("mmode-resv", region);
        }

        // 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;
            log::info!("Reserve initrd @ {:#x}..{:#x}", aligned_start, aligned_end);
            reserve_region(
                "initrd",
                PhysicalMemoryRegion {
                    base: aligned_start,
                    size,
                },
            );
        }

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

        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) };
            let initrd = self.initrd.init(data);

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

        Ok(())
    }

    // TODO boot hart ID may not be zero?
    unsafe fn init_platform(
        &'static self,
        hart_id: u32,
        queue_index: usize,
        is_bsp: bool,
    ) -> Result<(), Error> {
        let per_cpu = PerCpuData {
            tmp_t0: 0,
            umode_sp: 0,
            smode_sp: 0,

            bootstrap: is_bsp,
            queue_index,
        };
        Cpu::init_local(Some(hart_id), per_cpu);

        assert_eq!(Cpu::local().id(), hart_id);

        exception::init_smode_exceptions();

        if is_bsp {
            call_init_array();

            atomic::compiler_fence(Ordering::SeqCst);

            libk::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);

            unflatten_device_tree(dt);

            if let Some(machine_compatible) = machine_compatible {
                Self::apply_board_workarounds(machine_compatible);
            }

            if let Err(error) = Self::setup_clock_timebase(dt) {
                log::error!("Could not setup clock timebase from device tree: {error:?}");
            }
            if let Err(error) = Self::setup_chosen_stdout(dt) {
                log::error!("chosen-stdout setup error: {error:?}");
            } else {
                libk::debug::disable_early_sinks();
            }

            if let Some(machine) = machine_name {
                log::info!("Running on {machine:?}");
                MACHINE_NAME.init(machine.into());
            }
            log::info!("Boot arguments: {bootargs:?}");
            log::info!("Initializing riscv64 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()?;
        }

        // TODO more granular control over how U-mode pages are accessed from S-mode
        SSTATUS.modify(SSTATUS::SUM::SET);

        // Setup the timer
        SIE.modify(SIE::SSIE::SET + SIE::SEIE::SET);
        timer::init_hart(is_bsp);

        Ok(())
    }

    fn apply_board_workarounds(compatible: &str) {
        #[allow(clippy::single_match)]
        match compatible {
            _ => (),
        }
    }

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

    fn setup_clock_timebase(dt: &'static DeviceTree) -> Result<(), Error> {
        let cpus = dt.root().child("cpus").ok_or(Error::DoesNotExist)?;
        let timebase_frequency = cpus
            .prop_cell("timebase-frequency", 1)
            .ok_or(Error::DoesNotExist)?;
        timer::FREQUENCY.store(timebase_frequency, Ordering::Release);
        log::info!("System timer frequency: {}", Hertz(timebase_frequency));
        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 {
            log::info!("Probe chosen stdout: {:?}", node.name());
            node.force_init(InitSequence::Early)?;
        }

        // No stdout
        Ok(())
    }
}
