// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::collections::BTreeMap;
use std::ffi::c_void;
use std::ptr;

use anyhow::bail;
use dashmap::DashMap;
use nohash_hasher::BuildNoHashHasher;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use tracing::{debug, error};

#[derive(Debug, Clone, Copy)]
struct RegionEntry {
    end: usize,
    phys: usize,
}

#[derive(Debug)]
struct RegionMap {
    map: BTreeMap<usize, RegionEntry>,
}

impl RegionMap {
    fn new() -> Self {
        RegionMap {
            map: BTreeMap::new(),
        }
    }

    fn insert(&mut self, start: usize, end: usize, phys: usize) -> Option<RegionEntry> {
        self.map.insert(start, RegionEntry { end, phys })
    }

    fn query(&self, virt: usize) -> Option<usize> {
        let entry = self.map.range(..=virt).next_back();
        if let Some((&start, &RegionEntry { end, phys })) = entry {
            if virt < end || virt == end && start == end {
                let offset = virt - start;
                return Some(phys + offset);
            }
        }

        None
    }
}

static HANDLE_MAP: Lazy<RwLock<RegionMap>> = Lazy::new(|| RwLock::new(RegionMap::new()));
static REQ_ID_VHANDLE_MAP: Lazy<DashMap<u64, usize, BuildNoHashHasher<u64>>> =
    Lazy::new(|| DashMap::with_hasher(BuildNoHashHasher::default()));

pub fn is_vhandle_valid(virt: *mut c_void) -> bool {
    HANDLE_MAP.read().query(virt as usize).is_some()
}

pub fn handle_insert(virt: *mut c_void, phys: *mut c_void, size: usize) -> anyhow::Result<()> {
    if phys.is_null() {
        error!("[Virt] Failed to insert: phys is null");
        bail!("Failed to insert: phys is null");
    }

    let start = virt as usize;
    let end = start.wrapping_add(size);
    let mut map = HANDLE_MAP.write();
    if map.insert(start, end, phys as usize).is_some() {
        debug!(
            "[Virt] Updated mapping virt base {:p} -> (end: {:>#x}, phys {:p})",
            virt, end, phys
        );
    } else {
        debug!(
            "[Virt] Inserted mapping virt base {:p} -> (end: {:>#x}, phys {:p})",
            virt, end, phys
        );
    }

    Ok(())
}

pub fn handle_map(virt: *mut c_void) -> anyhow::Result<*mut c_void> {
    if virt.is_null() {
        debug!("[Virt] Mapping NULL handle, may be default API param");
        return Ok(ptr::null_mut());
    }

    if let Some(phys) = HANDLE_MAP.read().query(virt as usize) {
        debug!(
            "[Virt] Mapped virt {:p} -> phys {:p}",
            virt, phys as *mut c_void
        );
        Ok(phys as *mut c_void)
    } else {
        bail!("Error mapping virt: {:p} not found in any region", virt);
    }
}

pub fn req_id_vhandle_insert(request_id: u64, virt: *mut c_void) -> anyhow::Result<()> {
    let value = virt as usize;
    REQ_ID_VHANDLE_MAP.insert(request_id, value);
    debug!(
        "[Virt] Inserted request_id {} -> virt {:p}",
        request_id, virt
    );

    Ok(())
}

pub fn req_id_vhandle_map(request_id: u64) -> anyhow::Result<*mut c_void> {
    if let Some(v) = REQ_ID_VHANDLE_MAP.get(&request_id) {
        let virt = *v.value() as *mut c_void;
        debug!("[Virt] Mapped request_id {} -> virt {:p}", request_id, virt);

        Ok(virt)
    } else {
        error!("[Virt] Error mapping request_id {}: not found", request_id);
        bail!("Error mapping request_id {}: not found", request_id);
    }
}