#ifndef LIBSPF2_MOCK_H
#define LIBSPF2_MOCK_H
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <string>
#include <spf2/spf.h>

// Control variables to adjust behavior of mocked functions
extern "C" {
    static SPF_errcode_t spf_set_ipv4_result = SPF_E_SUCCESS;
    static SPF_errcode_t spf_set_helo_result = SPF_E_SUCCESS;
    static int spf_set_from_result;
    static SPF_response_t* spf_mock_response = nullptr;
    static SPF_errcode_t spf_request_query_result = SPF_E_SUCCESS;
    
    // New control variables for SPF_server mocking
    static SPF_server_t* spf_mock_server = nullptr;
    static bool spf_server_new_should_fail = false;

    SPF_server_t* SPF_server_new(SPF_server_dnstype_t dnstype,int debug)
    {
        if (spf_server_new_should_fail) {
            return nullptr;
        }
        return spf_mock_server ? spf_mock_server : new SPF_server_t();
    }

    void SPF_server_free(SPF_server_t* server) {
        delete server;
    }

    SPF_request_t* SPF_request_new(SPF_server_t *server) {
        return new SPF_request_t(); // Simply allocate and return new request
    }

    void SPF_request_free(SPF_request_t* request) {
        delete request; // Deallocate request
    }

    SPF_errcode_t SPF_request_set_ipv4_str(SPF_request_t* request, const char* ip) {
        return spf_set_ipv4_result; // Use controllable result
    }

    SPF_errcode_t SPF_request_set_helo_dom(SPF_request_t* request, const char* helo) {
        return spf_set_helo_result; // Use controllable result
    }

    int SPF_request_set_env_from(SPF_request_t* request, const char* from) {
        return spf_set_from_result; // Use controllable result
    }

    SPF_errcode_t SPF_request_query_mailfrom(SPF_request_t* request, SPF_response_t** response) {
        *response = spf_mock_response;
        return spf_request_query_result;
    }

    SPF_result_t SPF_response_result(SPF_response_t* response) {
        return response->result; // Return the result
    }

    void SPF_response_free(SPF_response_t* response) {
        delete response; // Deallocate response
    }
}

// Functions to set return values for mocked functions
void SetSpfMockReturnIPv4Value(SPF_errcode_t value) {
    spf_set_ipv4_result = value;
}

void SetSpfMockReturnHeloValue(SPF_errcode_t value) {
    spf_set_helo_result = value;
}

void SetSpfMockReturnFromValue(int value) {
    spf_set_from_result = value;
}

void SetSpfRequestQueryResult(SPF_errcode_t value) {
    spf_request_query_result = value;
}

void SetSpfMockResponse(SPF_response_t* mockResponse) {
    spf_mock_response = mockResponse;
}

// New functions for SPF_server mocking
void SetSpfMockServer(SPF_server_t* mockServer) {
    spf_mock_server = mockServer;
}

void SetSpfServerNewShouldFail(bool shouldFail) {
    spf_server_new_should_fail = shouldFail;
}

#endif // LIBSPF2_MOCK_H