#include "radar_protocol.h"
#include "target_tracker.h"

#include <math.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define ARRAY_SIZE(array) (sizeof(array) / sizeof((array)[0]))

static unsigned int tests_run;

static void fail(const char *file, int line, const char *expression)
{
    (void)fprintf(stderr, "%s:%d: assertion failed: %s\n", file, line, expression);
    exit(EXIT_FAILURE);
}

#define ASSERT_TRUE(expression)                                                \
    do {                                                                       \
        if (!(expression)) {                                                   \
            fail(__FILE__, __LINE__, #expression);                             \
        }                                                                      \
    } while (false)

#define ASSERT_EQ_INT(expected, actual)                                        \
    do {                                                                       \
        const long long expected_value = (long long)(expected);                \
        const long long actual_value = (long long)(actual);                    \
        if (expected_value != actual_value) {                                  \
            fail(__FILE__, __LINE__, #expected " == " #actual);               \
        }                                                                      \
    } while (false)

#define ASSERT_NEAR(expected, actual, tolerance)                               \
    do {                                                                       \
        const double difference = fabs((double)(expected) - (double)(actual)); \
        if (difference > (double)(tolerance)) {                                \
            fail(__FILE__, __LINE__, #expected " ~= " #actual);               \
        }                                                                      \
    } while (false)

#define RUN_TEST(function)                                                     \
    do {                                                                       \
        function();                                                            \
        ++tests_run;                                                           \
        (void)printf("ok %u - %s\n", tests_run, #function);                   \
    } while (false)

static uint16_t encode_sign_magnitude(int16_t value)
{
    if (value >= 0) {
        return (uint16_t)value | UINT16_C(0x8000);
    }
    return (uint16_t)(-(int32_t)value);
}

static void put_u16_le(uint8_t *destination, uint16_t value)
{
    destination[0] = (uint8_t)(value & UINT16_C(0x00FF));
    destination[1] = (uint8_t)(value >> 8U);
}

static void make_frame(uint8_t frame[RD03D_FRAME_SIZE],
                       const rd03d_target_t targets[RD03D_MAX_TARGETS])
{
    size_t index;

    memset(frame, 0, RD03D_FRAME_SIZE);
    memcpy(frame, rd03d_frame_header, RD03D_FRAME_HEADER_SIZE);
    memcpy(frame + RD03D_FRAME_SIZE - RD03D_FRAME_TAIL_SIZE,
           rd03d_frame_tail,
           RD03D_FRAME_TAIL_SIZE);

    for (index = 0U; index < RD03D_MAX_TARGETS; ++index) {
        uint8_t *payload;

        if (!targets[index].valid) {
            continue;
        }
        payload = frame + RD03D_FRAME_HEADER_SIZE +
                  index * RD03D_TARGET_DATA_SIZE;
        put_u16_le(payload, encode_sign_magnitude(targets[index].x_mm));
        put_u16_le(payload + 2U, encode_sign_magnitude(targets[index].y_mm));
        put_u16_le(payload + 4U,
                   encode_sign_magnitude(targets[index].speed_cm_s));
        put_u16_le(payload + 6U, targets[index].resolution_mm);
    }
}

static rd03d_target_t detection(int16_t x,
                                int16_t y,
                                int16_t speed,
                                uint16_t resolution)
{
    const rd03d_target_t result = {
        .x_mm = x,
        .y_mm = y,
        .speed_cm_s = speed,
        .resolution_mm = resolution,
        .valid = true,
    };
    return result;
}

static void test_sign_magnitude_boundaries(void)
{
    ASSERT_EQ_INT(-782, rd03d_decode_sign_magnitude(0x0EU, 0x03U));
    ASSERT_EQ_INT(1713, rd03d_decode_sign_magnitude(0xB1U, 0x86U));
    ASSERT_EQ_INT(-16, rd03d_decode_sign_magnitude(0x10U, 0x00U));
    ASSERT_EQ_INT(0, rd03d_decode_sign_magnitude(0x00U, 0x00U));
    ASSERT_EQ_INT(0, rd03d_decode_sign_magnitude(0x00U, 0x80U));
    ASSERT_EQ_INT(-32767, rd03d_decode_sign_magnitude(0xFFU, 0x7FU));
    ASSERT_EQ_INT(32767, rd03d_decode_sign_magnitude(0xFFU, 0xFFU));
}

static void test_decode_complete_frame_and_resolution(void)
{
    const rd03d_target_t input[RD03D_MAX_TARGETS] = {
        {.x_mm = -782, .y_mm = 1713, .speed_cm_s = -16,
         .resolution_mm = 360U, .valid = true},
        {.x_mm = 32767, .y_mm = -32767, .speed_cm_s = 72,
         .resolution_mm = 65535U, .valid = true},
        {0},
    };
    uint8_t bytes[RD03D_FRAME_SIZE];
    rd03d_frame_t frame;

    make_frame(bytes, input);
    ASSERT_TRUE(rd03d_decode_frame(bytes, &frame));
    ASSERT_TRUE(frame.targets[0].valid);
    ASSERT_EQ_INT(-782, frame.targets[0].x_mm);
    ASSERT_EQ_INT(1713, frame.targets[0].y_mm);
    ASSERT_EQ_INT(-16, frame.targets[0].speed_cm_s);
    ASSERT_EQ_INT(360, frame.targets[0].resolution_mm);
    ASSERT_TRUE(frame.targets[1].valid);
    ASSERT_EQ_INT(32767, frame.targets[1].x_mm);
    ASSERT_EQ_INT(-32767, frame.targets[1].y_mm);
    ASSERT_EQ_INT(65535, frame.targets[1].resolution_mm);
    ASSERT_TRUE(!frame.targets[2].valid);

    bytes[RD03D_FRAME_SIZE - 1U] ^= 0x01U;
    ASSERT_TRUE(!rd03d_decode_frame(bytes, &frame));
}

typedef struct {
    rd03d_frame_t frames[8];
    size_t count;
} frame_capture_t;

static void capture_frame(const rd03d_frame_t *frame, void *context)
{
    frame_capture_t *capture = context;

    ASSERT_TRUE(capture->count < ARRAY_SIZE(capture->frames));
    capture->frames[capture->count++] = *frame;
}

static void test_parser_all_split_points(void)
{
    const rd03d_target_t targets[RD03D_MAX_TARGETS] = {
        {.x_mm = -1200, .y_mm = 3456, .speed_cm_s = 88,
         .resolution_mm = 240U, .valid = true},
        {0},
        {0},
    };
    uint8_t bytes[RD03D_FRAME_SIZE];
    size_t split;

    make_frame(bytes, targets);
    for (split = 0U; split <= RD03D_FRAME_SIZE; ++split) {
        rd03d_parser_t parser;
        frame_capture_t capture = {0};
        size_t decoded_count;

        rd03d_parser_init(&parser);
        decoded_count = rd03d_parser_feed(&parser,
                                          bytes,
                                          split,
                                          capture_frame,
                                          &capture);
        decoded_count += rd03d_parser_feed(&parser,
                                           bytes + split,
                                           RD03D_FRAME_SIZE - split,
                                           capture_frame,
                                           &capture);
        ASSERT_EQ_INT(1, decoded_count);
        ASSERT_EQ_INT(1, capture.count);
        ASSERT_EQ_INT(-1200, capture.frames[0].targets[0].x_mm);
        ASSERT_EQ_INT(1, parser.frames_decoded);
        ASSERT_EQ_INT(0, parser.malformed_frames);
    }
}

static void test_parser_noise_corruption_and_resynchronisation(void)
{
    const rd03d_target_t first_targets[RD03D_MAX_TARGETS] = {
        {.x_mm = -321, .y_mm = 2000, .speed_cm_s = 10,
         .resolution_mm = 120U, .valid = true},
        {0},
        {0},
    };
    const rd03d_target_t second_targets[RD03D_MAX_TARGETS] = {
        {.x_mm = 777, .y_mm = 2500, .speed_cm_s = -25,
         .resolution_mm = 480U, .valid = true},
        {0},
        {0},
    };
    const uint8_t noise[] = {0x00U, 0xAAU, 0xAAU, 0xFFU, 0x02U, 0x99U};
    uint8_t bad[RD03D_FRAME_SIZE];
    uint8_t good[RD03D_FRAME_SIZE];
    uint8_t stream[sizeof(noise) + RD03D_FRAME_SIZE * 2U];
    rd03d_parser_t parser;
    frame_capture_t capture = {0};
    size_t offset = 0U;

    make_frame(bad, first_targets);
    make_frame(good, second_targets);
    bad[RD03D_FRAME_SIZE - 2U] = 0x54U;
    memcpy(stream + offset, noise, sizeof(noise));
    offset += sizeof(noise);
    memcpy(stream + offset, bad, sizeof(bad));
    offset += sizeof(bad);
    memcpy(stream + offset, good, sizeof(good));
    offset += sizeof(good);

    rd03d_parser_init(&parser);
    ASSERT_EQ_INT(1,
                  rd03d_parser_feed(&parser,
                                    stream,
                                    offset,
                                    capture_frame,
                                    &capture));
    ASSERT_EQ_INT(1, capture.count);
    ASSERT_EQ_INT(777, capture.frames[0].targets[0].x_mm);
    ASSERT_EQ_INT(-25, capture.frames[0].targets[0].speed_cm_s);
    ASSERT_EQ_INT(1, parser.malformed_frames);
    ASSERT_EQ_INT(offset, parser.bytes_received);
}

static void test_parser_recovers_from_dropped_byte_and_embedded_header(void)
{
    rd03d_target_t corrupt_targets[RD03D_MAX_TARGETS] = {
        {.x_mm = -100, .y_mm = 1000, .speed_cm_s = 4,
         .resolution_mm = 0U, .valid = true},
        {0},
        {0},
    };
    const rd03d_target_t good_targets[RD03D_MAX_TARGETS] = {
        {.x_mm = 901, .y_mm = 3333, .speed_cm_s = 42,
         .resolution_mm = 360U, .valid = true},
        {0},
        {0},
    };
    uint8_t corrupt[RD03D_FRAME_SIZE];
    uint8_t good[RD03D_FRAME_SIZE];
    uint8_t stream[RD03D_FRAME_SIZE * 2U - 1U];
    rd03d_parser_t parser;
    frame_capture_t capture = {0};

    make_frame(corrupt, corrupt_targets);
    make_frame(good, good_targets);
    /* A false header in the damaged payload must not prevent later recovery. */
    memcpy(corrupt + 12U, rd03d_frame_header, RD03D_FRAME_HEADER_SIZE);
    memcpy(stream, corrupt, 20U);
    memcpy(stream + 20U, corrupt + 21U, RD03D_FRAME_SIZE - 21U);
    memcpy(stream + RD03D_FRAME_SIZE - 1U, good, RD03D_FRAME_SIZE);

    rd03d_parser_init(&parser);
    ASSERT_EQ_INT(1,
                  rd03d_parser_feed(&parser,
                                    stream,
                                    sizeof(stream),
                                    capture_frame,
                                    &capture));
    ASSERT_EQ_INT(1, capture.count);
    ASSERT_EQ_INT(901, capture.frames[0].targets[0].x_mm);
    ASSERT_TRUE(parser.malformed_frames >= 1U);
}

static void test_parser_multiple_frames_per_chunk(void)
{
    rd03d_target_t targets[RD03D_MAX_TARGETS] = {
        {.x_mm = 1, .y_mm = 2, .speed_cm_s = 3,
         .resolution_mm = 4U, .valid = true},
        {0},
        {0},
    };
    uint8_t stream[RD03D_FRAME_SIZE * 3U];
    rd03d_parser_t parser;
    frame_capture_t capture = {0};
    size_t index;

    for (index = 0U; index < 3U; ++index) {
        targets[0].x_mm = (int16_t)(100 + (int16_t)index);
        make_frame(stream + index * RD03D_FRAME_SIZE, targets);
    }
    rd03d_parser_init(&parser);
    ASSERT_EQ_INT(3,
                  rd03d_parser_feed(&parser,
                                    stream,
                                    sizeof(stream),
                                    capture_frame,
                                    &capture));
    ASSERT_EQ_INT(3, capture.count);
    ASSERT_EQ_INT(100, capture.frames[0].targets[0].x_mm);
    ASSERT_EQ_INT(102, capture.frames[2].targets[0].x_mm);
}

static target_tracker_config_t deterministic_tracker_config(void)
{
    target_tracker_config_t config;

    target_tracker_default_config(&config);
    config.position_ema_alpha = 0.5F;
    config.speed_ema_alpha = 0.5F;
    config.motion_ema_alpha = 1.0F;
    config.max_association_distance_mm = 1500.0F;
    config.max_missed_frames = 2U;
    return config;
}

static const target_snapshot_t *find_snapshot(const target_snapshot_t *snapshots,
                                               size_t count,
                                               uint32_t id)
{
    size_t index;

    for (index = 0U; index < count; ++index) {
        if (snapshots[index].id == id) {
            return &snapshots[index];
        }
    }
    return NULL;
}

static void test_tracker_stable_ids_slot_reorder_and_ema(void)
{
    target_tracker_t tracker;
    target_tracker_config_t config = deterministic_tracker_config();
    rd03d_target_t frame_one[] = {
        {.x_mm = -1000, .y_mm = 2000, .speed_cm_s = 20,
         .resolution_mm = 120U, .valid = true},
        {.x_mm = 1000, .y_mm = 2200, .speed_cm_s = -30,
         .resolution_mm = 240U, .valid = true},
    };
    rd03d_target_t frame_two[] = {
        {.x_mm = 1100, .y_mm = 2200, .speed_cm_s = -10,
         .resolution_mm = 360U, .valid = true},
        {.x_mm = -800, .y_mm = 2000, .speed_cm_s = 40,
         .resolution_mm = 480U, .valid = true},
    };
    target_snapshot_t snapshots[TARGET_TRACKER_CAPACITY];
    uint32_t left_id;
    uint32_t right_id;
    size_t count;
    const target_snapshot_t *left;
    const target_snapshot_t *right;

    target_tracker_init(&tracker, &config);
    ASSERT_EQ_INT(2,
                  target_tracker_update(&tracker,
                                        frame_one,
                                        ARRAY_SIZE(frame_one)));
    count = target_tracker_snapshot(&tracker, snapshots, ARRAY_SIZE(snapshots));
    ASSERT_EQ_INT(2, count);
    left_id = snapshots[0].x_mm < 0.0F ? snapshots[0].id : snapshots[1].id;
    right_id = snapshots[0].x_mm > 0.0F ? snapshots[0].id : snapshots[1].id;

    ASSERT_EQ_INT(2,
                  target_tracker_update(&tracker,
                                        frame_two,
                                        ARRAY_SIZE(frame_two)));
    count = target_tracker_snapshot(&tracker, snapshots, ARRAY_SIZE(snapshots));
    left = find_snapshot(snapshots, count, left_id);
    right = find_snapshot(snapshots, count, right_id);
    ASSERT_TRUE(left != NULL);
    ASSERT_TRUE(right != NULL);
    ASSERT_NEAR(-900.0, left->x_mm, 0.001);
    ASSERT_NEAR(30.0, left->speed_cm_s, 0.001);
    ASSERT_EQ_INT(480, left->resolution_mm);
    ASSERT_NEAR(1050.0, right->x_mm, 0.001);
    ASSERT_NEAR(-20.0, right->speed_cm_s, 0.001);
    ASSERT_EQ_INT(360, right->resolution_mm);
    ASSERT_EQ_INT(2, left->age_frames);
    ASSERT_EQ_INT(0, left->last_seen_age_frames);
    ASSERT_TRUE(left->visible && left->observed_this_frame);
}

static void test_tracker_does_not_merge_close_targets(void)
{
    target_tracker_t tracker;
    rd03d_target_t detections[] = {
        {.x_mm = 0, .y_mm = 2000, .speed_cm_s = 1,
         .resolution_mm = 120U, .valid = true},
        {.x_mm = 500, .y_mm = 2000, .speed_cm_s = 2,
         .resolution_mm = 120U, .valid = true},
        {.x_mm = 900, .y_mm = 2000, .speed_cm_s = 3,
         .resolution_mm = 120U, .valid = true},
    };
    target_snapshot_t snapshots[TARGET_TRACKER_CAPACITY];

    target_tracker_init(&tracker, NULL);
    ASSERT_EQ_INT(3,
                  target_tracker_update(&tracker,
                                        detections,
                                        ARRAY_SIZE(detections)));
    ASSERT_EQ_INT(3,
                  target_tracker_snapshot(&tracker,
                                          snapshots,
                                          ARRAY_SIZE(snapshots)));
    ASSERT_TRUE(snapshots[0].id != snapshots[1].id);
    ASSERT_TRUE(snapshots[1].id != snapshots[2].id);
}

static void test_tracker_frame_age_removal_and_reacquisition(void)
{
    target_tracker_t tracker;
    target_tracker_config_t config = deterministic_tracker_config();
    const rd03d_target_t first = detection(100, 2000, 10, 120U);
    target_snapshot_t snapshot;
    uint32_t original_id;

    target_tracker_init(&tracker, &config);
    ASSERT_EQ_INT(1, target_tracker_update(&tracker, &first, 1U));
    ASSERT_EQ_INT(1, target_tracker_snapshot(&tracker, &snapshot, 1U));
    original_id = snapshot.id;

    ASSERT_EQ_INT(1, target_tracker_update(&tracker, NULL, 0U));
    ASSERT_EQ_INT(1, target_tracker_snapshot(&tracker, &snapshot, 1U));
    ASSERT_EQ_INT(1, snapshot.last_seen_age_frames);
    ASSERT_TRUE(snapshot.visible);
    ASSERT_TRUE(!snapshot.observed_this_frame);

    ASSERT_EQ_INT(1, target_tracker_update(&tracker, NULL, 0U));
    ASSERT_EQ_INT(1, target_tracker_snapshot(&tracker, &snapshot, 1U));
    ASSERT_EQ_INT(2, snapshot.last_seen_age_frames);
    ASSERT_EQ_INT(0, target_tracker_update(&tracker, NULL, 0U));
    ASSERT_EQ_INT(0, target_tracker_snapshot(&tracker, &snapshot, 1U));

    ASSERT_EQ_INT(1, target_tracker_update(&tracker, &first, 1U));
    ASSERT_EQ_INT(1, target_tracker_snapshot(&tracker, &snapshot, 1U));
    ASSERT_TRUE(snapshot.id != original_id);
}

static void test_tracker_prediction_preserves_ids_through_crossing(void)
{
    target_tracker_t tracker;
    target_tracker_config_t config = deterministic_tracker_config();
    rd03d_target_t detections[2];
    target_snapshot_t snapshots[2];
    uint32_t left_to_right_id;
    uint32_t right_to_left_id;
    size_t count;
    const target_snapshot_t *left_to_right;
    const target_snapshot_t *right_to_left;

    config.position_ema_alpha = 1.0F;
    target_tracker_init(&tracker, &config);
    detections[0] = detection(-600, 2000, 5, 100U);
    detections[1] = detection(600, 2000, -5, 100U);
    ASSERT_EQ_INT(2, target_tracker_update(&tracker, detections, 2U));
    count = target_tracker_snapshot(&tracker, snapshots, ARRAY_SIZE(snapshots));
    left_to_right_id = snapshots[0].x_mm < 0.0F ? snapshots[0].id : snapshots[1].id;
    right_to_left_id = snapshots[0].x_mm > 0.0F ? snapshots[0].id : snapshots[1].id;

    detections[0] = detection(-200, 2000, 5, 100U);
    detections[1] = detection(200, 2000, -5, 100U);
    ASSERT_EQ_INT(2, target_tracker_update(&tracker, detections, 2U));

    /* Swap input slots as the two physical trajectories cross. */
    detections[0] = detection(-200, 2000, -5, 100U);
    detections[1] = detection(200, 2000, 5, 100U);
    ASSERT_EQ_INT(2, target_tracker_update(&tracker, detections, 2U));
    count = target_tracker_snapshot(&tracker, snapshots, ARRAY_SIZE(snapshots));
    left_to_right = find_snapshot(snapshots, count, left_to_right_id);
    right_to_left = find_snapshot(snapshots, count, right_to_left_id);
    ASSERT_TRUE(left_to_right != NULL && right_to_left != NULL);
    ASSERT_NEAR(200.0, left_to_right->x_mm, 0.001);
    ASSERT_NEAR(-200.0, right_to_left->x_mm, 0.001);
}

static void test_tracker_replaces_missing_track_when_capacity_is_full(void)
{
    target_tracker_t tracker;
    target_tracker_config_t config = deterministic_tracker_config();
    rd03d_target_t detections[3] = {
        {.x_mm = -2000, .y_mm = 2000, .speed_cm_s = 1,
         .resolution_mm = 100U, .valid = true},
        {.x_mm = 0, .y_mm = 2000, .speed_cm_s = 2,
         .resolution_mm = 100U, .valid = true},
        {.x_mm = 2000, .y_mm = 2000, .speed_cm_s = 3,
         .resolution_mm = 100U, .valid = true},
    };
    target_snapshot_t snapshots[3];
    uint32_t departed_id = 0U;
    size_t count;

    config.position_ema_alpha = 1.0F;
    target_tracker_init(&tracker, &config);
    ASSERT_EQ_INT(3, target_tracker_update(&tracker, detections, 3U));
    count = target_tracker_snapshot(&tracker, snapshots, ARRAY_SIZE(snapshots));
    for (size_t i = 0U; i < count; ++i) {
        if (snapshots[i].x_mm == 2000.0F) {
            departed_id = snapshots[i].id;
        }
    }
    ASSERT_TRUE(departed_id != 0U);

    detections[2] = detection(6000, 2000, 4, 120U);
    ASSERT_EQ_INT(3, target_tracker_update(&tracker, detections, 3U));
    count = target_tracker_snapshot(&tracker, snapshots, ARRAY_SIZE(snapshots));
    ASSERT_EQ_INT(3, count);
    ASSERT_TRUE(find_snapshot(snapshots, count, departed_id) == NULL);

    bool found_new_target = false;
    for (size_t i = 0U; i < count; ++i) {
        if (snapshots[i].x_mm == 6000.0F) {
            found_new_target = true;
            ASSERT_TRUE(snapshots[i].id != departed_id);
        }
    }
    ASSERT_TRUE(found_new_target);
}

static void test_tracker_invalid_detection_and_reset(void)
{
    target_tracker_t tracker;
    const rd03d_target_t detections[] = {
        {.x_mm = 100, .y_mm = 200, .valid = false},
        {.x_mm = 300, .y_mm = 400, .speed_cm_s = 5,
         .resolution_mm = 10U, .valid = true},
    };
    target_snapshot_t snapshot;

    target_tracker_init(&tracker, NULL);
    ASSERT_EQ_INT(1,
                  target_tracker_update(&tracker,
                                        detections,
                                        ARRAY_SIZE(detections)));
    ASSERT_EQ_INT(1, target_tracker_snapshot(&tracker, &snapshot, 1U));
    ASSERT_EQ_INT(1, snapshot.id);
    target_tracker_reset(&tracker);
    ASSERT_EQ_INT(0, target_tracker_snapshot(&tracker, &snapshot, 1U));
    ASSERT_EQ_INT(1,
                  target_tracker_update(&tracker,
                                        detections,
                                        ARRAY_SIZE(detections)));
    ASSERT_EQ_INT(1, target_tracker_snapshot(&tracker, &snapshot, 1U));
    ASSERT_EQ_INT(1, snapshot.id);
}

int main(void)
{
    RUN_TEST(test_sign_magnitude_boundaries);
    RUN_TEST(test_decode_complete_frame_and_resolution);
    RUN_TEST(test_parser_all_split_points);
    RUN_TEST(test_parser_noise_corruption_and_resynchronisation);
    RUN_TEST(test_parser_recovers_from_dropped_byte_and_embedded_header);
    RUN_TEST(test_parser_multiple_frames_per_chunk);
    RUN_TEST(test_tracker_stable_ids_slot_reorder_and_ema);
    RUN_TEST(test_tracker_does_not_merge_close_targets);
    RUN_TEST(test_tracker_frame_age_removal_and_reacquisition);
    RUN_TEST(test_tracker_prediction_preserves_ids_through_crossing);
    RUN_TEST(test_tracker_replaces_missing_track_when_capacity_is_full);
    RUN_TEST(test_tracker_invalid_detection_and_reset);

    (void)printf("1..%u\n", tests_run);
    return EXIT_SUCCESS;
}