// SPDX-License-Identifier: GPL-2.0
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

use core::marker::PhantomData;

use kernel::{
    device,
    io::{
        poll::read_poll_timeout,
        register::Array,
        Io,
        Mmio, //
    },
    prelude::*,
    time::Delta, //
};

use crate::{
    driver::Bar0,
    falcon::{
        hal::LoadMethod,
        Falcon,
        FalconBromParams,
        FalconEngine,
        FalconModSelAlgo,
        PFalcon2Registers,
        PeregrineCoreSelect, //
    },
    regs,
};

use super::FalconHal;

fn select_core_ga102(pfalcon2: Mmio<'_, PFalcon2Registers>) -> Result {
    let bcr_ctrl = pfalcon2.read(regs::NV_PRISCV_RISCV_BCR_CTRL);
    if bcr_ctrl.core_select() != PeregrineCoreSelect::Falcon {
        pfalcon2.write_reg(
            regs::NV_PRISCV_RISCV_BCR_CTRL::zeroed().with_core_select(PeregrineCoreSelect::Falcon),
        );

        // TIMEOUT: falcon core should take less than 10ms to report being enabled.
        read_poll_timeout(
            || Ok(pfalcon2.read(regs::NV_PRISCV_RISCV_BCR_CTRL)),
            |r| r.valid(),
            Delta::ZERO,
            Delta::from_millis(10),
        )?;
    }

    Ok(())
}

fn signature_reg_fuse_version_ga102(
    dev: &device::Device,
    bar: Bar0<'_>,
    engine_id_mask: u16,
    ucode_id: u8,
) -> Result<u32> {
    // Each engine has 16 ucode version registers numbered from 1 to 16.
    let ucode_idx = match usize::from(ucode_id) {
        ucode_id @ 1..=regs::NV_FUSE_OPT_FPF_SIZE => ucode_id - 1,
        _ => {
            dev_err!(dev, "invalid ucode id {:#x}\n", ucode_id);
            return Err(EINVAL);
        }
    };

    // `ucode_idx` is guaranteed to be in the range [0..15], making the `read` calls provable valid
    // at build-time.
    let reg_fuse_version: u16 = if engine_id_mask & 0x0001 != 0 {
        bar.read(regs::NV_FUSE_OPT_FPF_SEC2_UCODE1_VERSION::at(ucode_idx))
            .data()
    } else if engine_id_mask & 0x0004 != 0 {
        bar.read(regs::NV_FUSE_OPT_FPF_NVDEC_UCODE1_VERSION::at(ucode_idx))
            .data()
    } else if engine_id_mask & 0x0400 != 0 {
        bar.read(regs::NV_FUSE_OPT_FPF_GSP_UCODE1_VERSION::at(ucode_idx))
            .data()
    } else {
        dev_err!(dev, "unexpected engine_id_mask {:#x}\n", engine_id_mask);
        return Err(EINVAL);
    };

    // TODO[NUMM]: replace with `last_set_bit` once it lands.
    Ok(u16::BITS - reg_fuse_version.leading_zeros())
}

fn program_brom_ga102(pfalcon2: Mmio<'_, PFalcon2Registers>, params: &FalconBromParams) {
    pfalcon2.write(
        Array::at(0),
        regs::NV_PFALCON2_FALCON_BROM_PARAADDR::zeroed().with_value(params.pkc_data_offset),
    );
    pfalcon2.write_reg(
        regs::NV_PFALCON2_FALCON_BROM_ENGIDMASK::zeroed()
            .with_value(u32::from(params.engine_id_mask)),
    );
    pfalcon2.write_reg(
        regs::NV_PFALCON2_FALCON_BROM_CURR_UCODE_ID::zeroed().with_ucode_id(params.ucode_id),
    );
    pfalcon2
        .write_reg(regs::NV_PFALCON2_FALCON_MOD_SEL::zeroed().with_algo(FalconModSelAlgo::Rsa3k));
}

pub(super) struct Ga102<E: FalconEngine>(PhantomData<E>);

impl<E: FalconEngine> Ga102<E> {
    pub(super) fn new() -> Self {
        Self(PhantomData)
    }
}

impl<E: FalconEngine> FalconHal<E> for Ga102<E> {
    fn select_core(&self, falcon: &Falcon<'_, E>) -> Result {
        select_core_ga102(falcon.pfalcon2)
    }

    fn signature_reg_fuse_version(
        &self,
        falcon: &Falcon<'_, E>,
        engine_id_mask: u16,
        ucode_id: u8,
    ) -> Result<u32> {
        signature_reg_fuse_version_ga102(falcon.dev, falcon.bar, engine_id_mask, ucode_id)
    }

    fn program_brom(&self, falcon: &Falcon<'_, E>, params: &FalconBromParams) {
        program_brom_ga102(falcon.pfalcon2, params);
    }

    fn is_riscv_active(&self, falcon: &Falcon<'_, E>) -> bool {
        falcon
            .pfalcon2
            .read(regs::NV_PRISCV_RISCV_CPUCTL)
            .active_stat()
    }

    fn is_riscv_halted(&self, falcon: &Falcon<'_, E>) -> Result<bool> {
        Ok(falcon.pfalcon2.read(regs::NV_PRISCV_RISCV_CPUCTL).halted())
    }

    fn reset_wait_mem_scrubbing(&self, falcon: &Falcon<'_, E>) -> Result {
        // TIMEOUT: memory scrubbing should complete in less than 20ms.
        read_poll_timeout(
            || Ok(falcon.pfalcon.read(regs::NV_PFALCON_FALCON_HWCFG2)),
            |r| r.mem_scrubbing_done(),
            Delta::ZERO,
            Delta::from_millis(20),
        )
        .map(|_| ())
    }

    fn reset_eng(&self, falcon: &Falcon<'_, E>) -> Result {
        let _ = falcon.pfalcon.read(regs::NV_PFALCON_FALCON_HWCFG2);

        // According to OpenRM's `kflcnPreResetWait_GA102` documentation, HW sometimes does not set
        // RESET_READY so a non-failing timeout is used.
        let _ = read_poll_timeout(
            || Ok(falcon.pfalcon.read(regs::NV_PFALCON_FALCON_HWCFG2)),
            |r| r.reset_ready(),
            Delta::ZERO,
            Delta::from_micros(150),
        );

        regs::NV_PFALCON_FALCON_ENGINE::reset_engine(falcon.pfalcon);
        self.reset_wait_mem_scrubbing(falcon)?;

        Ok(())
    }

    fn load_method(&self) -> LoadMethod {
        LoadMethod::Dma
    }

    fn host_routed_causes(
        &self,
        falcon: &Falcon<'_, E>,
        latched: regs::NV_PFALCON_FALCON_IRQSTAT,
    ) -> regs::NV_PFALCON_FALCON_IRQSTAT {
        let pfalcon2 = falcon.pfalcon2;
        let mask = pfalcon2.read(regs::ga102::NV_PRISCV_RISCV_IRQMASK).value();
        let dest = pfalcon2.read(regs::ga102::NV_PRISCV_RISCV_IRQDEST).value();

        regs::NV_PFALCON_FALCON_IRQSTAT::from(latched.into_raw() & mask & dest)
    }

    fn retrigger(&self, falcon: &Falcon<'_, E>) {
        super::tu102::retrigger_ga100(falcon);
    }
}