use super::*;
use crate::{bitvec, expr, variable};

impl State {
    /// Generate a mock state for an ARM-32 state.
    pub fn mock_arm32() -> State {
        State::new(
            &Tid::new("mock_fn"),
            &variable!("sp:4"),
            &CallingConvention::mock_arm32(),
        )
    }

    /// Generate a mock state for an x64 state.
    pub fn mock_x64(tid_name: &str) -> State {
        State::new(
            &Tid::new(tid_name),
            &variable!("RSP:8"),
            &CallingConvention::mock_x64(),
        )
    }
}

/// Mock an abstract ID representing the stack.
fn mock_stack_id() -> AbstractIdentifier {
    AbstractIdentifier::from_var(Tid::new("mock_fn"), &variable!("sp:4"))
}

/// Mock an abstract ID of a stack parameter
fn mock_stack_param_id(offset: i64, size: u64) -> AbstractIdentifier {
    AbstractIdentifier::new(
        Tid::new("mock_fn"),
        AbstractLocation::from_stack_position(
            mock_stack_id().unwrap_register(),
            offset,
            ByteSize::new(size),
        ),
    )
}

#[test]
fn test_new() {
    let state = State::mock_arm32();
    // Test the generated stack
    assert_eq!(&state.stack_id, &mock_stack_id());
    assert_eq!(state.stack.iter().len(), 0);
    // Assert that the register values are as expected
    assert_eq!(state.register.len(), 7); // 6 parameter register plus stack pointer
    assert_eq!(
        state.get_register(&variable!("sp:4")),
        DataDomain::from_target(mock_stack_id(), bitvec!("0x0:4").into())
    );
    // Check the generated tracked IDs
    assert_eq!(state.tracked_ids.len(), 6);
    for (id, access_pattern) in state.tracked_ids.iter() {
        assert_eq!(
            state.get_register(id.unwrap_register()),
            DataDomain::from_target(
                id.clone(),
                bitvec!(format!("0:{}", id.unwrap_register().size)).into()
            )
        );
        assert_eq!(access_pattern, &AccessPattern::new());
    }
}

#[test]
fn test_store_and_load_from_stack() {
    let mut state = State::mock_arm32();
    let address = DataDomain::from_target(mock_stack_id(), bitvec!("-4:4").into());
    let value: DataDomain<BitvectorDomain> = bitvec!("0x0:4").into();
    // write and load a value to the current stack frame
    state.write_value(address.clone(), value.clone());
    assert_eq!(state.stack.iter().len(), 1);
    assert_eq!(
        state.stack.get(bitvec!("-4:4"), ByteSize::new(4)),
        value.clone()
    );
    assert_eq!(state.load_value(address, ByteSize::new(4), None), value);
    // Load a parameter register and check that the parameter gets generated
    let address = DataDomain::from_target(mock_stack_id(), bitvec!("0x4:4").into());
    let stack_param_id = mock_stack_param_id(4, 4);
    let stack_param = DataDomain::from_target(stack_param_id.clone(), bitvec!("0x0:4").into());
    assert_eq!(state.tracked_ids.iter().len(), 6);
    assert_eq!(
        state.load_value(address.clone(), ByteSize::new(4), None),
        stack_param
    );
    assert_eq!(state.tracked_ids.iter().len(), 7);
    assert_eq!(
        state
            .tracked_ids
            .get(&stack_param_id)
            .unwrap()
            .is_accessed(),
        false
    ); // The load method does not set access flags.
}

#[test]
fn test_load_unsized_from_stack() {
    let mut state = State::mock_arm32();
    // Load an existing stack param (generated by a sized load at the same address)
    let address = DataDomain::from_target(mock_stack_id(), bitvec!("0x0:4").into());
    let stack_param_id = mock_stack_param_id(0, 4);
    let stack_param = DataDomain::from_target(stack_param_id.clone(), bitvec!("0x0:4").into());
    state.load_value(address, ByteSize::new(4), None);
    let unsized_load = state.load_unsized_value_from_stack(bitvec!("0x0:4").into());
    assert_eq!(unsized_load, stack_param);
    assert!(state.tracked_ids.get(&stack_param_id).is_some());
    // Load a non-existing stack param
    let stack_param_id = mock_stack_param_id(4, 1);
    let stack_param = DataDomain::from_target(stack_param_id.clone(), bitvec!("0x0:1").into());
    let unsized_load = state.load_unsized_value_from_stack(bitvec!("0x4:4"));
    assert_eq!(unsized_load, stack_param);
    assert!(state.tracked_ids.get(&stack_param_id).is_some());
    // Unsized load from the current stack frame
    let unsized_load = state.load_unsized_value_from_stack(bitvec!("-4:4"));
    assert_eq!(unsized_load, DataDomain::new_top(ByteSize::new(1)));
}

#[test]
fn test_eval() {
    let mut state = State::mock_arm32();
    // Test the eval method
    let expr = expr!("sp:4 + 42:4");
    assert_eq!(
        state.eval(&expr),
        DataDomain::from_target(mock_stack_id(), bitvec!("42:4").into())
    );
    // Test the eval_parameter_arg method
    let arg = Arg::from_var(variable!("sp:4"), None);
    assert_eq!(
        state.eval_parameter_arg(&arg),
        DataDomain::from_target(mock_stack_id(), bitvec!("0x0:4").into())
    );
}

#[test]
fn test_extern_symbol_handling() {
    let mut state = State::mock_arm32();
    let extern_symbol = ExternSymbol::mock_arm32("mock_symbol");
    let cconv = CallingConvention::mock_arm32();
    let call_tid = Tid::new("call_tid");
    let param_id = AbstractIdentifier::from_var(Tid::new("mock_fn"), &variable!("r0:4"));
    let return_val_id = AbstractIdentifier::from_var(Tid::new("call_tid"), &variable!("r0:4"));
    // Test extern symbol handling.
    state.handle_generic_extern_symbol(
        &call_tid,
        &extern_symbol,
        &cconv,
        &RuntimeMemoryImage::mock(),
    );
    assert_eq!(
        state
            .tracked_ids
            .get(&param_id)
            .unwrap()
            .is_mutably_dereferenced(),
        true
    );
    let return_val = state.get_register(&variable!("r0:4"));
    assert_eq!(return_val.get_relative_values().iter().len(), 2);
    assert_eq!(
        return_val.get_relative_values().get(&param_id).unwrap(),
        &BitvectorDomain::new_top(ByteSize::new(4))
    );
    assert_eq!(
        return_val.get_relative_values().get(&param_id).unwrap(),
        &BitvectorDomain::new_top(ByteSize::new(4))
    );
    assert_eq!(
        return_val
            .get_relative_values()
            .get(&return_val_id)
            .unwrap(),
        &bitvec!("0:4").into()
    );
}

#[test]
fn test_substitute_global_mem_address() {
    let mut state = State::mock_arm32();
    let global_memory = RuntimeMemoryImage::mock();

    // Test that addresses into non-writeable memory do not get substituted.
    let global_address: DataDomain<BitvectorDomain> = bitvec!("0x1000:4").into();
    let substituted_address =
        state.substitute_global_mem_address(global_address.clone(), &global_memory);
    assert_eq!(global_address, substituted_address);
    // Test substitution for addresses into writeable global memory.
    let global_address: DataDomain<BitvectorDomain> = bitvec!("0x2000:4").into();
    let substituted_address = state.substitute_global_mem_address(global_address, &global_memory);
    let expected_global_id = AbstractIdentifier::from_global_address(
        state.get_current_function_tid(),
        &bitvec!("0x2000:4"),
    );
    assert_eq!(
        state.tracked_ids.get(&expected_global_id),
        Some(&AccessPattern::new())
    );
    assert_eq!(
        substituted_address,
        DataDomain::from_target(expected_global_id, bitvec!("0x0:4").into())
    );
}
