#include "aiw_soc_peer.h"

#include <assert.h>
#include <stdio.h>
#include <string.h>

typedef struct {
    aiw_soc_capture capture;
    aiw_soc_inference inference;
    bool capture_ok;
    bool inference_ok;
    unsigned capture_calls;
    unsigned inference_calls;
    uint64_t last_minimum;
} fake_soc;

static bool fake_capture(void *context, const aiw_ticket *request,
                         aiw_soc_capture *capture)
{
    fake_soc *soc = context;
    assert(soc != NULL && request != NULL && capture != NULL);
    ++soc->capture_calls;
    soc->last_minimum = request->min_source_frame_seq;
    if (!soc->capture_ok) return false;
    *capture = soc->capture;
    return true;
}

static bool fake_infer(void *context, const aiw_soc_capture *capture,
                       uint16_t channels, aiw_soc_inference *inference)
{
    fake_soc *soc = context;
    assert(soc != NULL && capture != NULL && inference != NULL);
    assert(channels == 61u);
    ++soc->inference_calls;
    if (!soc->inference_ok) return false;
    *inference = soc->inference;
    return true;
}

static size_t encode_request(uint8_t *packet, uint32_t ticket_id,
                             uint64_t minimum_frame_seq)
{
    const aiw_ticket request = {
        .session_id = UINT64_C(0x123456789ABCDEF0),
        .ticket_id = ticket_id,
        .calibration_id = 77u,
        .channels = 61u,
        .min_source_frame_seq = minimum_frame_seq
    };
    size_t length = 0u;
    assert(aiw_encode_ticket(&request, packet,
                             AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES +
                                 AIW_CRC_BYTES,
                             &length));
    return length;
}

static void init_peer(aiw_soc_peer *peer, fake_soc *soc)
{
    assert(aiw_soc_peer_init(peer, UINT64_C(0x123456789ABCDEF0), 77u, 61u,
                             fake_capture, fake_infer, soc));
}

static void test_fresh_frame_and_nail_mask(void)
{
    aiw_soc_peer peer = {0};
    fake_soc soc = {0};
    aiw_result result = {0};
    uint8_t request[AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES];
    uint8_t response[AIW_MAX_PACKET_BYTES];
    size_t response_length = 0u;
    const int image_token = 1;

    soc.capture_ok = true;
    soc.inference_ok = true;
    soc.capture = (aiw_soc_capture){
        .frame_handle = &image_token,
        .capture_baseline_seq = 0u,
        .source_frame_seq = 1u,
        .exposure_started_after_request = true
    };
    soc.inference.hand_present = true;
    soc.inference.nail_channel_mask[4] = 1u;
    soc.inference.pwm[4] = 127u;
    soc.inference.pwm[5] = 255u; /* Must be blanked outside the nail mask. */

    init_peer(&peer, &soc);
    const size_t request_length = encode_request(request, 1u, 1u);
    assert(aiw_soc_peer_process(&peer, request, request_length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_RESPONSE_ENCODED);
    assert(aiw_decode_result(response, response_length, &result));
    assert(result.session_id == peer.session_id && result.ticket_id == 1u);
    assert(result.capture_baseline_seq == 0u && result.source_frame_seq == 1u);
    assert(result.status_flags == AIW_RESULT_FLAG_HAND_PRESENT);
    assert(result.pwm[4] == 127u && result.pwm[5] == 0u);
    assert(soc.last_minimum == 1u && soc.capture_calls == 1u);
}

static void test_stale_cached_and_replayed_requests_fail_closed(void)
{
    aiw_soc_peer peer = {0};
    fake_soc soc = {0};
    uint8_t request[AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES];
    uint8_t response[AIW_MAX_PACKET_BYTES];
    size_t response_length = 99u;
    const int image_token = 1;

    soc.capture_ok = true;
    soc.inference_ok = true;
    soc.capture = (aiw_soc_capture){
        .frame_handle = &image_token,
        .capture_baseline_seq = 0u,
        .source_frame_seq = 1u,
        .exposure_started_after_request = false
    };
    init_peer(&peer, &soc);
    size_t length = encode_request(request, 1u, 1u);
    assert(aiw_soc_peer_process(&peer, request, length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_STALE_FRAME);
    assert(response_length == 0u && peer.last_source_frame_seq == 0u);
    assert(aiw_soc_peer_process(&peer, request, length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_REPLAYED_TICKET);
    assert(soc.capture_calls == 1u);

    soc.capture.exposure_started_after_request = true;
    length = encode_request(request, 2u, 1u);
    assert(aiw_soc_peer_process(&peer, request, length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_RESPONSE_ENCODED);
    assert(peer.last_source_frame_seq == 1u);

    soc.capture.source_frame_seq = 1u; /* Cached sequence cannot satisfy next ticket. */
    soc.capture.capture_baseline_seq = 0u;
    length = encode_request(request, 3u, 2u);
    assert(aiw_soc_peer_process(&peer, request, length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_STALE_FRAME);
    assert(response_length == 0u && peer.last_source_frame_seq == 1u);
}

static void test_mismatched_session_and_corrupt_packet_are_rejected(void)
{
    aiw_soc_peer peer = {0};
    fake_soc soc = {0};
    aiw_ticket request_value = {
        .session_id = UINT64_C(0x1020304050607080),
        .ticket_id = 1u,
        .calibration_id = 77u,
        .channels = 61u,
        .min_source_frame_seq = 1u
    };
    uint8_t request[AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES];
    uint8_t response[AIW_MAX_PACKET_BYTES];
    size_t length = 0u, response_length = 0u;

    init_peer(&peer, &soc);
    assert(aiw_encode_ticket(&request_value, request, sizeof(request), &length));
    assert(aiw_soc_peer_process(&peer, request, length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_SESSION_MISMATCH);
    request[10] ^= 0x01u;
    assert(aiw_soc_peer_process(&peer, request, length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_INVALID_REQUEST);
    assert(soc.capture_calls == 0u && peer.rejected_requests == 2u);
}

static void test_model_failure_returns_zero_output(void)
{
    aiw_soc_peer peer = {0};
    fake_soc soc = {0};
    aiw_result result = {0};
    const int image_token = 1;
    uint8_t request[AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES];
    uint8_t response[AIW_MAX_PACKET_BYTES];
    size_t response_length = 0u;

    soc.capture_ok = true;
    soc.inference_ok = false;
    soc.capture = (aiw_soc_capture){
        .frame_handle = &image_token,
        .capture_baseline_seq = 10u,
        .source_frame_seq = 11u,
        .exposure_started_after_request = true
    };
    init_peer(&peer, &soc);
    const size_t request_length = encode_request(request, 1u, 11u);
    assert(aiw_soc_peer_process(&peer, request, request_length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_RESPONSE_ENCODED);
    assert(aiw_decode_result(response, response_length, &result));
    assert(result.status_flags == 0u && peer.inference_failures == 1u);
    for (uint16_t i = 0u; i < result.channels; ++i) assert(result.pwm[i] == 0u);
}

static void test_capture_failure_consumes_ticket_and_sends_nothing(void)
{
    aiw_soc_peer peer = {0};
    fake_soc soc = {0};
    const int image_token = 1;
    uint8_t request[AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES];
    uint8_t response[AIW_MAX_PACKET_BYTES];
    size_t response_length = 17u;

    init_peer(&peer, &soc);
    size_t request_length = encode_request(request, 1u, 1u);
    assert(aiw_soc_peer_process(&peer, request, request_length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_CAPTURE_FAILED);
    assert(response_length == 0u && peer.capture_failures == 1u);
    assert(peer.last_ticket_id == 1u && soc.capture_calls == 1u);

    soc.capture_ok = true;
    soc.inference_ok = true;
    soc.capture = (aiw_soc_capture){
        .frame_handle = &image_token,
        .capture_baseline_seq = 0u,
        .source_frame_seq = 1u,
        .exposure_started_after_request = true
    };
    assert(aiw_soc_peer_process(&peer, request, request_length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_REPLAYED_TICKET);
    request_length = encode_request(request, 2u, 1u);
    assert(aiw_soc_peer_process(&peer, request, request_length, response,
                                sizeof(response), &response_length) ==
           AIW_SOC_PEER_RESPONSE_ENCODED);
}

static void test_small_response_buffer_never_gets_a_partial_packet(void)
{
    aiw_soc_peer peer = {0};
    fake_soc soc = {0};
    const int image_token = 1;
    uint8_t request[AIW_HEADER_BYTES + AIW_TICKET_PAYLOAD_BYTES + AIW_CRC_BYTES];
    uint8_t tiny_response[1] = {0xA5u};
    size_t response_length = 17u;

    soc.capture_ok = true;
    soc.inference_ok = true;
    soc.capture = (aiw_soc_capture){
        .frame_handle = &image_token,
        .capture_baseline_seq = 0u,
        .source_frame_seq = 1u,
        .exposure_started_after_request = true
    };
    init_peer(&peer, &soc);
    const size_t request_length = encode_request(request, 1u, 1u);
    assert(aiw_soc_peer_process(&peer, request, request_length, tiny_response,
                                sizeof(tiny_response), &response_length) ==
           AIW_SOC_PEER_RESPONSE_TOO_LARGE);
    assert(response_length == 0u && tiny_response[0] == 0xA5u);
}

int main(void)
{
    test_fresh_frame_and_nail_mask();
    test_stale_cached_and_replayed_requests_fail_closed();
    test_mismatched_session_and_corrupt_packet_are_rejected();
    test_model_failure_returns_zero_output();
    test_capture_failure_consumes_ticket_and_sends_nothing();
    test_small_response_buffer_never_gets_a_partial_packet();
    puts("PASS: AIW SoC peer requires fresh frames, rejects replay, masks non-nail PWM and fails dark without partial results.");
    return 0;
}
