use super::super::Runtime;
use super::super::contract::*;
use crate::wasm_runtime::runtime::RuntimeConfig;
use crate::wasm_runtime::tests::TestSetup;
use crate::wasm_runtime::{ContractExecError, RuntimeInnerError};
use freenet_stdlib::prelude::*;
use std::time::Instant;
use tracing::info;

const TEST_CONTRACT_METERING: &str = "test_contract_metering";

const HIGH_ITERATIONS: u64 = 5_000_000_000;

const TIMEOUT_ITERATIONS: u64 = 100_000_000_000;

#[derive(Debug, Default, serde::Serialize, serde::Deserialize)]
struct TestConditions {
    pub iterations: u64,
}

#[tokio::test(flavor = "multi_thread")]
async fn validate_state_metering() -> Result<(), Box<dyn std::error::Error>> {
    let TestSetup {
        contract_store,
        delegate_store,
        secrets_store,
        contract_key,
        temp_dir,
    } = super::setup_test_contract(TEST_CONTRACT_METERING).await?;

    let config = RuntimeConfig {
        max_execution_seconds: 5.0,
        cpu_cycles_per_second: Some(1_000_000), // Lower limit to force gas error
        safety_margin: 0.1,
        enable_metering: true,
        ..Default::default()
    };

    let mut runtime =
        Runtime::build_with_config(contract_store, delegate_store, secrets_store, false, config)
            .unwrap();

    let test_conditions = TestConditions {
        iterations: HIGH_ITERATIONS,
    };
    let state = WrappedState::new(serde_json::to_vec(&test_conditions)?);
    let time = Instant::now();

    let result = runtime.validate_state(
        &contract_key,
        &Parameters::from([].as_ref()),
        &state,
        &Default::default(),
    );

    let duration = time.elapsed().as_secs_f64();
    info!("Duration: {duration:.2}s");

    assert!(duration < 5.0, "Should not timeout");
    assert!(
        matches!(
            result.as_ref().err().map(|e| e.deref()),
            Some(RuntimeInnerError::ContractExecError(
                ContractExecError::OutOfGas
            ))
        ),
        "Should fail with gas error"
    );
    std::mem::drop(temp_dir);
    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_update_state_metering() -> Result<(), Box<dyn std::error::Error>> {
    let TestSetup {
        contract_store,
        delegate_store,
        secrets_store,
        contract_key,
        temp_dir,
    } = super::setup_test_contract(TEST_CONTRACT_METERING).await?;

    let config = RuntimeConfig {
        max_execution_seconds: 5.0,
        cpu_cycles_per_second: Some(2_000_000),
        safety_margin: 0.1,
        enable_metering: true,
        ..Default::default()
    };

    let mut runtime =
        Runtime::build_with_config(contract_store, delegate_store, secrets_store, false, config)
            .unwrap();

    let test_conditions = TestConditions {
        iterations: HIGH_ITERATIONS,
    };
    let state = WrappedState::new(serde_json::to_vec(&test_conditions)?);
    let time = Instant::now();

    let result = runtime.update_state(
        &contract_key,
        &Parameters::from([].as_ref()),
        &state,
        &[StateDelta::from([].as_ref()).into()],
    );

    let duration = time.elapsed().as_secs_f64();
    info!("Duration: {duration:.2}s");

    assert!(duration < 5.0, "Should not timeout");
    assert!(
        matches!(
            result.as_ref().err().map(|e| e.deref()),
            Some(RuntimeInnerError::ContractExecError(
                ContractExecError::OutOfGas
            ))
        ),
        "Should fail with gas error"
    );

    std::mem::drop(temp_dir);
    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_summarize_state_metering() -> Result<(), Box<dyn std::error::Error>> {
    let TestSetup {
        contract_store,
        delegate_store,
        secrets_store,
        contract_key,
        temp_dir,
    } = super::setup_test_contract(TEST_CONTRACT_METERING).await?;

    let config = RuntimeConfig {
        max_execution_seconds: 5.0,
        cpu_cycles_per_second: Some(3_000_000),
        safety_margin: 0.1,
        enable_metering: true,
        ..Default::default()
    };

    let mut runtime =
        Runtime::build_with_config(contract_store, delegate_store, secrets_store, false, config)
            .unwrap();

    let test_conditions = TestConditions {
        iterations: HIGH_ITERATIONS,
    };
    let state = WrappedState::new(serde_json::to_vec(&test_conditions)?);
    let time = Instant::now();

    let result = runtime.summarize_state(&contract_key, &Parameters::from([].as_ref()), &state);

    let duration = time.elapsed().as_secs_f64();
    info!("Duration: {duration:.2}s");

    assert!(duration < 5.0, "Should not timeout");
    assert!(
        matches!(
            result.as_ref().err().map(|e| e.deref()),
            Some(RuntimeInnerError::ContractExecError(
                ContractExecError::OutOfGas
            ))
        ),
        "Should fail with gas error"
    );

    std::mem::drop(temp_dir);
    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_get_state_delta_metering() -> Result<(), Box<dyn std::error::Error>> {
    let TestSetup {
        contract_store,
        delegate_store,
        secrets_store,
        contract_key,
        temp_dir,
    } = super::setup_test_contract(TEST_CONTRACT_METERING).await?;

    let config = RuntimeConfig {
        max_execution_seconds: 5.0,
        cpu_cycles_per_second: Some(4_000_000),
        safety_margin: 0.1,
        enable_metering: true,
        ..Default::default()
    };

    let mut runtime =
        Runtime::build_with_config(contract_store, delegate_store, secrets_store, false, config)
            .unwrap();

    let test_conditions = TestConditions {
        iterations: HIGH_ITERATIONS,
    };
    let state = WrappedState::new(serde_json::to_vec(&test_conditions)?);
    let time = Instant::now();

    let result = runtime.get_state_delta(
        &contract_key,
        &Parameters::from([].as_ref()),
        &state,
        &StateSummary::from([].as_ref()),
    );

    let duration = time.elapsed().as_secs_f64();
    info!("Duration: {duration:.2}s");

    assert!(duration < 5.0, "Should not timeout");
    assert!(
        matches!(
            result.as_ref().err().map(|e| e.deref()),
            Some(RuntimeInnerError::ContractExecError(
                ContractExecError::OutOfGas
            ))
        ),
        "Should fail with gas error"
    );

    std::mem::drop(temp_dir);
    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_timeout_metering() -> Result<(), Box<dyn std::error::Error>> {
    let TestSetup {
        contract_store,
        delegate_store,
        secrets_store,
        contract_key,
        temp_dir,
    } = super::setup_test_contract(TEST_CONTRACT_METERING).await?;

    let config = RuntimeConfig {
        max_execution_seconds: 5.0,
        cpu_cycles_per_second: Some(u64::MAX),
        safety_margin: 0.1,
        enable_metering: true,
        ..Default::default()
    };

    let mut runtime =
        Runtime::build_with_config(contract_store, delegate_store, secrets_store, false, config)
            .unwrap();

    let test_conditions = TestConditions {
        iterations: TIMEOUT_ITERATIONS,
    };
    let state = WrappedState::new(serde_json::to_vec(&test_conditions)?);
    let time = Instant::now();

    let result = runtime.validate_state(
        &contract_key,
        &Parameters::from([].as_ref()),
        &state,
        &Default::default(),
    );

    let duration = time.elapsed().as_secs_f64();
    // Allow for some timing variance due to system load - should timeout around 5s but allow up to 7s
    assert!(
        duration < 7.0,
        "Took {duration:.2}s, should timeout before 7s (configured for 5s timeout)"
    );

    assert!(
        matches!(
            result.as_ref().err().map(|e| e.deref()),
            Some(RuntimeInnerError::ContractExecError(
                ContractExecError::MaxComputeTimeExceeded
            ))
        ),
        "Should fail with timeout error"
    );

    std::mem::drop(temp_dir);
    Ok(())
}
