// This file is part of Gear.

// Copyright (C) 2022-2024 Gear Technologies Inc.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Lazy pages support for runtime.

#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

use byteorder::{ByteOrder, LittleEndian};
use core::fmt;
use gear_core::{
    ids::ProgramId,
    memory::{HostPointer, Memory, MemoryInterval},
    pages::{GearPage, PageNumber, PageU32Size, WasmPage},
    program::MemoryInfix,
};
use gear_lazy_pages_common::{GlobalsAccessConfig, LazyPagesWeights, ProcessAccessError, Status};
use gear_runtime_interface::{gear_ri, LazyPagesProgramContext};
use sp_std::{mem, vec::Vec};

fn mprotect_lazy_pages(mem: &mut impl Memory, protect: bool) {
    if mem.get_buffer_host_addr().is_none() {
        return;
    }

    // Cannot panic, unless OS has some problems with pages protection.
    gear_ri::mprotect_lazy_pages(protect);
}

/// Try to enable and initialize lazy pages env
pub fn try_to_enable_lazy_pages(prefix: [u8; 32]) -> bool {
    gear_ri::init_lazy_pages(gear_lazy_pages_common::LazyPagesInitContext::new(prefix).into())
}

/// Protect and save storage keys for pages which has no data
pub fn init_for_program(
    mem: &mut impl Memory,
    program_id: ProgramId,
    memory_infix: MemoryInfix,
    stack_end: Option<WasmPage>,
    globals_config: GlobalsAccessConfig,
    weights: LazyPagesWeights,
) {
    let weights = [
        weights.signal_read,
        weights.signal_write,
        weights.signal_write_after_read,
        weights.host_func_read,
        weights.host_func_write,
        weights.host_func_write_after_read,
        weights.load_page_storage_data,
    ]
    .map(|w| w.one())
    .to_vec();

    let ctx = LazyPagesProgramContext {
        wasm_mem_addr: mem.get_buffer_host_addr(),
        wasm_mem_size: mem.size().raw(),
        stack_end: stack_end.map(|p| p.raw()),
        program_key: {
            let memory_infix = memory_infix.inner().to_le_bytes();

            [program_id.as_ref(), memory_infix.as_ref()].concat()
        },
        globals_config,
        weights,
    };

    // Cannot panic unless OS allocates buffer in not aligned by native page addr, or
    // something goes wrong with pages protection.
    gear_ri::init_lazy_pages_for_program(ctx);
}

/// Remove lazy-pages protection, returns wasm memory begin addr
pub fn remove_lazy_pages_prot(mem: &mut impl Memory) {
    mprotect_lazy_pages(mem, false);
}

/// Protect lazy-pages and set new wasm mem addr and size,
/// if they have been changed.
pub fn update_lazy_pages_and_protect_again(
    mem: &mut impl Memory,
    old_mem_addr: Option<HostPointer>,
    old_mem_size: WasmPage,
    new_mem_addr: HostPointer,
) {
    struct PointerDisplay(HostPointer);

    impl fmt::Debug for PointerDisplay {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "{:#x}", self.0)
        }
    }

    let changed_addr = if old_mem_addr
        .map(|addr| new_mem_addr != addr)
        .unwrap_or(true)
    {
        log::debug!(
            "backend executor has changed wasm mem buff: from {:?} to {:?}",
            old_mem_addr.map(PointerDisplay),
            new_mem_addr
        );

        Some(new_mem_addr)
    } else {
        None
    };

    let new_mem_size = mem.size();
    let changed_size = (new_mem_size > old_mem_size).then_some(new_mem_size.raw());

    if !matches!((changed_addr, changed_size), (None, None)) {
        gear_ri::change_wasm_memory_addr_and_size(changed_addr, changed_size)
    }

    mprotect_lazy_pages(mem, true);
}

/// Returns list of released pages numbers.
pub fn get_write_accessed_pages() -> Vec<GearPage> {
    gear_ri::write_accessed_pages()
        .into_iter()
        .map(|p| {
            GearPage::new(p)
                .unwrap_or_else(|_| unreachable!("Lazy pages backend returns wrong pages"))
        })
        .collect()
}

/// Returns lazy pages actual status.
pub fn get_status() -> Status {
    gear_ri::lazy_pages_status().0
}

fn serialize_mem_intervals(intervals: &[MemoryInterval]) -> Vec<u8> {
    let mut bytes = Vec::with_capacity(mem::size_of_val(intervals));
    for interval in intervals {
        bytes.extend_from_slice(&interval.to_bytes());
    }
    bytes
}

/// Pre-process memory access in syscalls in lazy-pages.
pub fn pre_process_memory_accesses(
    reads: &[MemoryInterval],
    writes: &[MemoryInterval],
    gas_counter: &mut u64,
) -> Result<(), ProcessAccessError> {
    let serialized_reads = serialize_mem_intervals(reads);
    let serialized_writes = serialize_mem_intervals(writes);

    let mut gas_bytes = [0u8; 8];
    LittleEndian::write_u64(&mut gas_bytes, *gas_counter);

    let res =
        gear_ri::pre_process_memory_accesses(&serialized_reads, &serialized_writes, &mut gas_bytes);

    *gas_counter = LittleEndian::read_u64(&gas_bytes);

    // if result can be converted to `ProcessAccessError` then it's an error
    if let Ok(err) = ProcessAccessError::try_from(res) {
        return Err(err);
    }
    Ok(())
}