#include "lamp_ui_wire_service.h"
#include "aiw_soc_ui_client.h"
#include "aiw_stream.h"

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

static unsigned checks;

#define CHECK(expr) do { \
    ++checks; \
    if (!(expr)) { \
        fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #expr); \
        exit(EXIT_FAILURE); \
    } \
} while (0)

static aiw_ui_ack dispatch(lamp_ui_wire_service *service, uint64_t session_id,
                           uint32_t sequence, uint8_t button, uint32_t now_ms) {
    aiw_message message = {0};
    aiw_ui_ack ack = {0}, decoded = {0};
    uint8_t packet[AIW_MAX_PACKET_BYTES];
    size_t length = 0u;
    message.type = AIW_MESSAGE_UI_COMMAND;
    message.value.ui_command.session_id = session_id;
    message.value.ui_command.sequence = sequence;
    message.value.ui_command.button = button;
    CHECK(lamp_ui_wire_service_handle(service, &message, now_ms, &ack));
    CHECK(aiw_encode_ui_ack(&ack, packet, sizeof(packet), &length));
    CHECK(aiw_decode_ui_ack(packet, length, &decoded));
    return decoded;
}

static void open_session(lamp_ui_wire_service *service, uint64_t session_id,
                         uint32_t now_ms, aiw_ui_ack_status expected) {
    aiw_message message = {0};
    aiw_ui_ack ack = {0}, decoded = {0};
    uint8_t packet[AIW_MAX_PACKET_BYTES];
    size_t length = 0u;
    message.type = AIW_MESSAGE_UI_SESSION_OPEN;
    message.value.ui_session_open.session_id = session_id;
    CHECK(lamp_ui_wire_service_handle(service, &message, now_ms, &ack));
    CHECK(ack.sequence == 0u && ack.status == expected);
    CHECK(aiw_encode_ui_ack(&ack, packet, sizeof(packet), &length));
    CHECK(aiw_decode_ui_ack(packet, length, &decoded));
    CHECK(decoded.session_id == session_id && decoded.status == expected);
}

static void setup_guard(lg_guard *guard, uint64_t session_id) {
    const lg_config config = {
        .channels = 2u,
        .max_run_ms = LG_MANUAL_MAX_RUN_MS,
        .ticket_lifetime_ms = 900000u,
        .max_poll_gap_ms = 900000u,
        .channel_input_mw = {100u, 100u},
        .fixed_input_mw = 100u,
        .input_budget_mw = 500u,
        .calibration_id = 7u
    };
    const lg_inputs safe_inputs = {true, true, true, true};
    CHECK(lg_init(guard, &config, session_id, 100u));
    lg_poll(guard, 100u, safe_inputs);
}

static aiw_ui_ack feed_soc_packet(aiw_stream *stream,
                                  lamp_ui_wire_service *service,
                                  const uint8_t *packet, size_t packet_length,
                                  uint32_t now_ms) {
    aiw_message message = {0};
    aiw_ui_ack ack = {0}, decoded = {0};
    uint8_t ack_packet[AIW_MAX_PACKET_BYTES];
    size_t ack_length = 0u;
    bool received = false;
    for (size_t i = 0u; i < packet_length; ++i) {
        if (aiw_stream_push_message(stream, packet[i], now_ms++, &message)) {
            CHECK(!received);
            received = true;
        }
    }
    CHECK(received);
    CHECK(message.type == AIW_MESSAGE_UI_SESSION_OPEN ||
          message.type == AIW_MESSAGE_UI_COMMAND);
    CHECK(lamp_ui_wire_service_handle(service, &message, now_ms, &ack));
    CHECK(aiw_encode_ui_ack(&ack, ack_packet, sizeof(ack_packet), &ack_length));
    CHECK(aiw_decode_ui_ack(ack_packet, ack_length, &decoded));
    return decoded;
}

int main(void) {
    const uint64_t session_id = UINT64_C(0x123456789ABCDEF0);
    const uint64_t other_session = UINT64_C(0x0FEDCBA987654321);
    lg_guard guard = {0};
    lamp_ui_wire_service service;
    setup_guard(&guard, session_id);
    CHECK(lamp_ui_wire_service_init(&service, &guard));

    CHECK(dispatch(&service, session_id, 1u, LUI_BUTTON_CANCEL, 110u).status ==
          AIW_UI_ACK_BAD_SESSION);
    open_session(&service, session_id, 120u, AIW_UI_ACK_SESSION_OPENED);
    CHECK(dispatch(&service, session_id, 1u, LUI_BUTTON_LONG_PRESS_MANUAL, 130u).status ==
          AIW_UI_ACK_ACCEPTED);
    aiw_ui_ack ack = dispatch(&service, session_id, 2u, LUI_BUTTON_MANUAL_60, 140u);
    CHECK(ack.status == AIW_UI_ACK_ACCEPTED);
    CHECK((ack.flags & AIW_UI_ACK_FLAG_RUNNING) != 0u && ack.remaining_ms == 60000u);
    const uint32_t start_ms = guard.run_started_ms;

    /* Reopening the same session is idempotent and preserves the command replay cache. */
    open_session(&service, session_id, 150u, AIW_UI_ACK_SESSION_OPENED);
    ack = dispatch(&service, session_id, 2u, LUI_BUTTON_MANUAL_60, 1000u);
    CHECK(ack.status == AIW_UI_ACK_ACCEPTED);
    CHECK((ack.flags & AIW_UI_ACK_FLAG_DUPLICATE) != 0u);
    CHECK(ack.remaining_ms == 59140u && guard.run_started_ms == start_ms);
    CHECK(dispatch(&service, session_id, 2u, LUI_BUTTON_CANCEL, 1010u).status ==
          AIW_UI_ACK_SEQUENCE_COLLISION);
    ack = dispatch(&service, other_session, 1u, LUI_BUTTON_CANCEL, 1015u);
    CHECK(ack.status == AIW_UI_ACK_BAD_SESSION && ack.selected_mode == LG_OFF &&
          ack.flags == 0u && ack.remaining_ms == 0u);

    open_session(&service, other_session, 1020u, AIW_UI_ACK_SESSION_REJECTED);
    ack = dispatch(&service, session_id, 3u, LUI_BUTTON_CANCEL, 1030u);
    CHECK(ack.status == AIW_UI_ACK_ACCEPTED &&
          (ack.flags & AIW_UI_ACK_FLAG_RUNNING) == 0u);
    open_session(&service, other_session, 1040u, AIW_UI_ACK_SESSION_REJECTED);
    CHECK(!guard.armed && !guard.output_enable);

    /* The actual target's deliberately unconfigured guard accepts the session,
       but refuses all timed output commands with NOT_READY. */
    lg_guard unconfigured_guard = {0};
    lamp_ui_wire_service unconfigured_service;
    const uint64_t boot_session = UINT64_C(0xA0B0C0D0E0F00112);
    CHECK(lamp_ui_wire_service_init(&unconfigured_service, &unconfigured_guard));
    open_session(&unconfigured_service, boot_session, 2000u,
                 AIW_UI_ACK_SESSION_OPENED);
    ack = dispatch(&unconfigured_service, boot_session, 1u,
                   LUI_BUTTON_LONG_PRESS_MANUAL, 2005u);
    CHECK(ack.status == AIW_UI_ACK_ACCEPTED);
    ack = dispatch(&unconfigured_service, boot_session, 2u,
                   LUI_BUTTON_MANUAL_60, 2010u);
    CHECK(ack.status == AIW_UI_ACK_NOT_READY);
    CHECK(!unconfigured_guard.armed && !unconfigured_guard.output_enable);

    /* Exercise the complete SoC client -> UART parser -> target service -> ACK path. */
    lg_guard round_trip_guard = {0};
    lamp_ui_wire_service round_trip_service;
    aiw_soc_ui_client client;
    aiw_stream stream;
    uint8_t packet[AIW_MAX_PACKET_BYTES], retry[AIW_MAX_PACKET_BYTES];
    size_t packet_length = 0u, retry_length = 0u;
    const uint64_t round_trip_session = UINT64_C(0x3141592653589793);
    setup_guard(&round_trip_guard, round_trip_session);
    CHECK(lamp_ui_wire_service_init(&round_trip_service, &round_trip_guard));
    CHECK(aiw_soc_ui_client_init(&client, round_trip_session));
    CHECK(aiw_stream_init(&stream, AIW_STREAM_DEFAULT_TIMEOUT_MS));
    CHECK(aiw_soc_ui_client_encode_session_open(&client, packet,
                                                sizeof(packet), &packet_length));
    ack = feed_soc_packet(&stream, &round_trip_service, packet, packet_length, 3000u);
    CHECK(aiw_soc_ui_client_accept_session_ack(&client, &ack));
    CHECK(client.session_open);
    CHECK(aiw_soc_ui_client_begin_command(&client, LUI_BUTTON_LONG_PRESS_MANUAL,
                                         packet, sizeof(packet), &packet_length));
    ack = feed_soc_packet(&stream, &round_trip_service, packet, packet_length, 3010u);
    CHECK(aiw_soc_ui_client_accept_command_ack(&client, &ack));
    CHECK(aiw_soc_ui_client_begin_command(&client, LUI_BUTTON_MANUAL_60,
                                         packet, sizeof(packet), &packet_length));
    const uint32_t expected_run_start = 3020u + (uint32_t)packet_length;
    const aiw_ui_ack dropped_ack = feed_soc_packet(&stream, &round_trip_service,
                                                   packet, packet_length, 3020u);
    CHECK(dropped_ack.status == AIW_UI_ACK_ACCEPTED);
    CHECK(aiw_soc_ui_client_retry_command(&client, retry, sizeof(retry), &retry_length));
    CHECK(retry_length == packet_length && memcmp(packet, retry, packet_length) == 0);
    ack = feed_soc_packet(&stream, &round_trip_service, retry, retry_length, 4000u);
    CHECK(ack.status == AIW_UI_ACK_ACCEPTED &&
          (ack.flags & AIW_UI_ACK_FLAG_DUPLICATE) != 0u);
    CHECK(aiw_soc_ui_client_accept_command_ack(&client, &ack));
    CHECK(!client.command_pending && round_trip_guard.run_started_ms == expected_run_start);
    CHECK(stream.accepted_frames == 4u);

    printf("PASS: lamp_ui_wire_service %u assertions\n", checks);
    return EXIT_SUCCESS;
}
