use crate::{
    module_factory::ModuleFactory,
    runtime::{ContractRuntime, ExecutionMode, SystemRuntime},
    RuntimeContext,
};
use fluentbase_types::{
    byteorder::{ByteOrder, LittleEndian},
    import_linker_v1_preview, Address, BytecodeOrHash, ExitCode, HashMap, B256,
};
use rwasm::{ExecutionEngine, ImportLinker, RwasmModule, StrategyDefinition, TrapCode};
use std::{cell::RefCell, mem::take, sync::Arc};

/// Finalized outcome of a single runtime invocation.
///
/// Values are reported in fuel units; gas conversion (if any) is handled by the caller.
#[derive(Default, Clone, Debug)]
pub struct ExecutionResult {
    /// Contract-defined exit status. Negative values map from TrapCode via ExitCode; zero indicates success.
    pub exit_code: i32,
    /// Total fuel consumed by the invocation (excludes refunded fuel).
    pub fuel_consumed: u64,
    /// Fuel refunded to the caller (negative values are not expected).
    pub fuel_refunded: i64,
    /// Raw output buffer produced by the callee; for nested calls it is moved into the parent's return_data.
    pub output: Vec<u8>,
    /// Return data propagated back to the parent on success paths of nested calls.
    pub return_data: Vec<u8>,
}

impl ExecutionResult {
    pub fn take_and_continue(&mut self, is_interrupted: bool) -> Self {
        let mut result = take(self);
        // We don't propagate output into an intermediary state
        if is_interrupted {
            self.output = take(&mut result.output);
            self.return_data = take(&mut result.return_data);
        }
        result
    }
}

/// Captures an intentional execution interruption that must be resumed by the root context.
#[derive(Debug, Default, Clone)]
pub struct ExecutionInterruption {
    /// Fuel spent up to the interruption point.
    pub fuel_consumed: u64,
    /// Fuel to refund to the caller at the interruption point.
    pub fuel_refunded: i64,
    /// Encoded interruption payload (e.g., delegated call parameters).
    pub return_data: Vec<u8>,
}

/// Result of running or resuming a runtime.
#[derive(Clone, Debug)]
pub enum RuntimeResult {
    /// Execution finished; contains the finalized result.
    Result(ExecutionResult),
    /// Execution yielded; contains data necessary to resume later.
    Interruption(ExecutionInterruption),
}

impl RuntimeResult {
    /// Unwraps the successful execution result; panics if this is an interruption.
    pub fn into_execution_result(self) -> ExecutionResult {
        match self {
            RuntimeResult::Result(result) => result,
            _ => unreachable!(),
        }
    }
}

pub trait RuntimeExecutor {
    /// Executes the entry function of the module determined by the current execution state.
    ///
    /// Returns either a finalized result.
    fn execute(&mut self, bytecode_or_hash: BytecodeOrHash, ctx: RuntimeContext)
        -> ExecutionResult;

    /// Resumes a previously interrupted runtime.
    ///
    /// `fuel16_ptr` optionally points to a 16-byte buffer where fuel consumption and refund are written back.
    fn resume(
        &mut self,
        call_id: u32,
        return_data: &[u8],
        fuel16_ptr: u32,
        fuel_consumed: u64,
        fuel_refunded: i64,
        exit_code: i32,
    ) -> ExecutionResult;

    /// Drop a runtime we don't need to resume anymore
    fn forget_runtime(&mut self, call_id: u32);

    /// Warm up the bytecode
    fn warmup(&mut self, bytecode: RwasmModule, hash: B256, address: Address);

    /// Resets the per-transaction call identifier counter and clears recoverable runtimes.
    ///
    /// Intended to be invoked at the beginning of a new transaction.
    fn reset_call_id_counter(&mut self);

    fn memory_read(
        &mut self,
        call_id: u32,
        offset: usize,
        buffer: &mut [u8],
    ) -> Result<(), TrapCode>;
}

pub struct ThreadLocalExecutor;

thread_local! {
    pub static LOCAL_RUNTIME_EXECUTOR: RefCell<RuntimeFactoryExecutor> = RefCell::new(RuntimeFactoryExecutor::new(import_linker_v1_preview()));
}

impl RuntimeExecutor for ThreadLocalExecutor {
    fn execute(
        &mut self,
        bytecode_or_hash: BytecodeOrHash,
        ctx: RuntimeContext,
    ) -> ExecutionResult {
        LOCAL_RUNTIME_EXECUTOR
            .with_borrow_mut(|runtime_executor| runtime_executor.execute(bytecode_or_hash, ctx))
    }

    fn resume(
        &mut self,
        call_id: u32,
        return_data: &[u8],
        fuel16_ptr: u32,
        fuel_consumed: u64,
        fuel_refunded: i64,
        exit_code: i32,
    ) -> ExecutionResult {
        LOCAL_RUNTIME_EXECUTOR.with_borrow_mut(|runtime_executor| {
            runtime_executor.resume(
                call_id,
                return_data,
                fuel16_ptr,
                fuel_consumed,
                fuel_refunded,
                exit_code,
            )
        })
    }

    fn forget_runtime(&mut self, call_id: u32) {
        LOCAL_RUNTIME_EXECUTOR
            .with_borrow_mut(|runtime_executor| runtime_executor.forget_runtime(call_id))
    }

    fn warmup(&mut self, bytecode: RwasmModule, hash: B256, address: Address) {
        LOCAL_RUNTIME_EXECUTOR
            .with_borrow_mut(|runtime_executor| runtime_executor.warmup(bytecode, hash, address))
    }

    fn reset_call_id_counter(&mut self) {
        LOCAL_RUNTIME_EXECUTOR
            .with_borrow_mut(|runtime_executor| runtime_executor.reset_call_id_counter())
    }

    fn memory_read(
        &mut self,
        call_id: u32,
        offset: usize,
        buffer: &mut [u8],
    ) -> Result<(), TrapCode> {
        LOCAL_RUNTIME_EXECUTOR.with_borrow_mut(|runtime_executor| {
            runtime_executor.memory_read(call_id, offset, buffer)
        })
    }
}

/// Returns a default runtime executor.
pub fn default_runtime_executor() -> impl RuntimeExecutor {
    ThreadLocalExecutor {}
}

pub struct RuntimeFactoryExecutor {
    /// A module factory
    pub module_factory: ModuleFactory,
    /// Suspended runtimes keyed by per-transaction call identifier.
    pub recoverable_runtimes: HashMap<u32, ExecutionMode>,
    /// An import linker
    pub import_linker: Arc<ImportLinker>,
    /// Monotonically increasing counter for assigning call identifiers.
    pub transaction_call_id_counter: u32,
}

impl RuntimeFactoryExecutor {
    pub fn new(import_linker: Arc<ImportLinker>) -> Self {
        Self {
            module_factory: ModuleFactory::new(),
            recoverable_runtimes: HashMap::new(),
            import_linker,
            transaction_call_id_counter: 1,
        }
    }

    /// Saves the current runtime instance for later resumption and returns its call identifier.
    pub fn try_remember_runtime(
        &mut self,
        runtime_result: RuntimeResult,
        runtime: ExecutionMode,
    ) -> ExecutionResult {
        let interruption = match runtime_result {
            RuntimeResult::Result(result) => {
                // Return result (there is no need to do anything else)
                return result;
            }
            RuntimeResult::Interruption(interruption) => interruption,
        };

        // Get current call_id before incrementing
        let call_id = self.transaction_call_id_counter;

        // Check if call_id would overflow i32 when cast (positive exit codes are reserved for call_id)
        if call_id > i32::MAX as u32 {
            return ExecutionResult {
                exit_code: ExitCode::UnknownError.into_i32(),
                fuel_consumed: interruption.fuel_consumed,
                fuel_refunded: interruption.fuel_refunded,
                output: vec![],
                return_data: vec![],
            };
        }

        // Increment counter for next call (safe since call_id <= i32::MAX < u32::MAX)
        self.transaction_call_id_counter += 1;
        let prev = self.recoverable_runtimes.insert(call_id, runtime);
        debug_assert!(prev.is_none());

        ExecutionResult {
            // We return `call_id` as exit code (it's safe because exit code can't be positive)
            exit_code: call_id as i32,
            // Forward info about consumed and refunded fuel (during the call)
            fuel_consumed: interruption.fuel_consumed,
            fuel_refunded: interruption.fuel_refunded,
            // The output we map into return data
            output: interruption.return_data,
            return_data: vec![],
        }
    }

    /// Consolidates the trap/result of an invocation into a RuntimeResult and updates accounting.
    ///
    /// When fuel_consumed_before_the_call is provided, computes precise fuel usage by diffing the
    /// store's remaining fuel. Returns either a finalized result or an interruption wrapper.
    fn handle_execution_result(
        &mut self,
        next_result: Result<(), TrapCode>,
        fuel_consumed: Option<u64>,
        ctx: &mut RuntimeContext,
    ) -> RuntimeResult {
        let mut execution_result = ctx
            .execution_result
            .take_and_continue(ctx.resumable_context.is_some());
        // There are two counters for fuel: opcode fuel counter; manually charged.
        // It's applied for execution runtimes where we don't know the final fuel consumed
        // till it's committed by Wasm runtime.
        // That is why we rewrite fuel here to check how much we've really spent based on the context information.
        if let Some(store_fuel_consumed) = fuel_consumed {
            execution_result.fuel_consumed = store_fuel_consumed;
        }

        // Fill the exit code in the execution result based on the next result:
        // - Ok - execution passed, exit code is 0 (Ok)
        // - InterruptionCalled - we don't know exit code since it's just an interruption
        // - Err - an execution trap code (halts execution)
        match next_result {
            Ok(_) => {
                // Don't write exit code here, because it's managed by host functions
            }
            Err(TrapCode::InterruptionCalled) => {
                // We don't set exit code here,
                // because exit code is used to represent identifier of call id
            }
            Err(err) => {
                execution_result.exit_code = ExitCode::from(err).into_i32();
            }
        }

        // If the next result is interruption
        if next_result == Err(TrapCode::InterruptionCalled) {
            let ExecutionResult {
                fuel_consumed,
                fuel_refunded,
                mut return_data,
                ..
            } = execution_result;
            // A case for normal interruption (not system runtime interruption), where we should
            // serialize the context we remembered inside the `exec.rs `
            // handler to pass into parent runtime.
            //
            // Safety: For system runtimes we don't save this
            if let Some(resumable_return_data) = ctx.take_resumable_context_serialized() {
                return_data = resumable_return_data;
            }
            return RuntimeResult::Interruption(ExecutionInterruption {
                fuel_consumed,
                fuel_refunded,
                return_data,
            });
        }

        RuntimeResult::Result(execution_result)
    }
}
impl RuntimeExecutor for RuntimeFactoryExecutor {
    fn execute(
        &mut self,
        bytecode_or_hash: BytecodeOrHash,
        ctx: RuntimeContext,
    ) -> ExecutionResult {
        let system_runtime_params = match &bytecode_or_hash {
            BytecodeOrHash::Bytecode { address, hash, .. } => {
                fluentbase_types::is_execute_using_system_runtime(address)
                    .then_some((*address, *hash))
            }
            BytecodeOrHash::Hash(_) => None,
        };

        // If we have a cached module, then use it, otherwise create a new one and cache
        let module = self.module_factory.get_module_or_init(bytecode_or_hash);

        // If there is no cached store, then construct a new one (slow)
        let fuel_limit_value = ctx.fuel_limit;
        let fuel_limit = Some(fuel_limit_value);

        let mut exec_mode = if let Some((address, code_hash)) = system_runtime_params {
            let consume_fuel = fluentbase_types::is_engine_metered_precompile(&address);
            let runtime = SystemRuntime::new(
                module,
                self.import_linker.clone(),
                code_hash,
                ctx,
                consume_fuel,
            );
            ExecutionMode::System(runtime)
        } else {
            let engine = ExecutionEngine::acquire_shared();
            // We always execute untrusted contracts with rWasm VM
            let strategy = StrategyDefinition::Rwasm { engine, module };
            let runtime =
                ContractRuntime::new(strategy, self.import_linker.clone(), ctx, fuel_limit);
            // This is an extraordinary case where we fail during resource init inside the entrypoint,
            // but there is nothing we can do here rather than just return the execution error.
            //
            // By default, start sections are not allowed for users, so users can't deploy contracts that
            // cause traps inside the entrypoint.
            //
            // Ideally, it should never happen.
            if let Some(trap_code) = runtime.as_ref().err() {
                return ExecutionResult {
                    exit_code: ExitCode::from(trap_code).into_i32(),
                    fuel_consumed: fuel_limit_value,
                    fuel_refunded: 0,
                    output: vec![],
                    return_data: vec![],
                };
            }
            ExecutionMode::Contract(runtime.unwrap())
        };

        // Execute program
        let result = exec_mode.execute();
        let fuel_consumed = exec_mode
            .remaining_fuel()
            .zip(fuel_limit)
            .map(|(remaining_fuel, store_fuel)| store_fuel - remaining_fuel);

        let runtime_result =
            self.handle_execution_result(result, fuel_consumed, exec_mode.context_mut());
        self.try_remember_runtime(runtime_result, exec_mode)
    }

    fn resume(
        &mut self,
        call_id: u32,
        return_data: &[u8],
        fuel16_ptr: u32,
        fuel_consumed: u64,
        fuel_refunded: i64,
        exit_code: i32,
    ) -> ExecutionResult {
        let Some(mut runtime) = self.recoverable_runtimes.remove(&call_id) else {
            unreachable!(
                "runtime: missing recoverable runtime for resume, this should never happen: call_id={}, fuel_consumed={}, exit_code={}",
                call_id, fuel_consumed, exit_code
            )
        };
        let mut fuel_remaining = runtime.remaining_fuel();
        let resume_inner = |runtime: &mut ExecutionMode| {
            // Copy return data into return data
            runtime.context_mut().execution_result.return_data = return_data.to_vec();
            if fuel16_ptr > 0 {
                let mut buffer = [0u8; 16];
                LittleEndian::write_u64(&mut buffer[..8], fuel_consumed);
                LittleEndian::write_i64(&mut buffer[8..], fuel_refunded);
                runtime.memory_write(fuel16_ptr as usize, &buffer)?;
            }
            runtime.resume(exit_code, fuel_consumed)
        };
        let result = resume_inner(&mut runtime);
        // We need to adjust the fuel limit because `fuel_consumed` should not be included into spent.
        if result != Err(TrapCode::OutOfFuel) {
            // Safety: We can safely unwrap here, because `OutOfFuel` check we did in `resume_inner` and the result is ok.
            fuel_remaining = fuel_remaining.map(|v| v.checked_sub(fuel_consumed).unwrap());
        }
        let fuel_consumed = runtime
            .remaining_fuel()
            .and_then(|remaining_fuel| Some(fuel_remaining? - remaining_fuel));
        let runtime_result =
            self.handle_execution_result(result, fuel_consumed, runtime.context_mut());
        self.try_remember_runtime(runtime_result, runtime)
    }

    fn forget_runtime(&mut self, call_id: u32) {
        _ = self.recoverable_runtimes.remove(&call_id);
    }

    fn warmup(&mut self, bytecode: RwasmModule, hash: B256, address: Address) {
        self.module_factory
            .get_module_or_init(BytecodeOrHash::Bytecode {
                bytecode,
                hash,
                address,
            });
    }

    fn reset_call_id_counter(&mut self) {
        // For each transaction we reset the `call_id` counter (used to track interruptions)
        self.transaction_call_id_counter = 1;
        // Clear recoverable runtimes, because they are no longer valid
        self.recoverable_runtimes.clear();
    }

    fn memory_read(
        &mut self,
        call_id: u32,
        offset: usize,
        buffer: &mut [u8],
    ) -> Result<(), TrapCode> {
        let runtime_ref = self.recoverable_runtimes.get_mut(&call_id).expect(
            "runtime: missing recoverable runtime for memory read, this should never happen",
        );
        runtime_ref.memory_read(offset, buffer)
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        executor::{ExecutionInterruption, RuntimeFactoryExecutor, RuntimeResult},
        runtime::{ContractRuntime, ExecutionMode},
        RuntimeContext,
    };
    use fluentbase_types::{import_linker_v1_preview, ExitCode};
    use rwasm::{ExecutionEngine, RwasmModule, StrategyDefinition};

    #[test]
    fn call_id_overflow() {
        let mut executor = RuntimeFactoryExecutor::new(import_linker_v1_preview());

        // Set counter to i32::MAX to trigger overflow on the next allocation
        executor.transaction_call_id_counter = i32::MAX as u32 + 1;

        let interruption = RuntimeResult::Interruption(ExecutionInterruption {
            fuel_consumed: 100,
            fuel_refunded: 0,
            return_data: vec![1, 2, 3],
        });

        let engine = ExecutionEngine::acquire_shared();
        let module = RwasmModule::default();
        let ctx = RuntimeContext::default();

        let strategy_runtime = ContractRuntime::new(
            StrategyDefinition::Rwasm { module, engine },
            executor.import_linker.clone(),
            ctx,
            None,
        )
        .unwrap();
        let runtime = ExecutionMode::Contract(strategy_runtime);

        // Try to allocate call_id - should fail with overflow
        let result = executor.try_remember_runtime(interruption, runtime);

        // Verify overflow error
        assert_eq!(result.exit_code, ExitCode::UnknownError.into_i32());
        assert_eq!(result.fuel_consumed, 100);
        assert_eq!(result.fuel_refunded, 0);
        assert!(result.output.is_empty());
    }
}
