#include "doctest.h"
#include <tactility/concurrent/thread.h>
#include <tactility/delay.h>
#include <tactility/device.h>
#include <tactility/driver.h>

#include <atomic>

static bool probe_present = false;
static int probe_driver_start_called = 0;
static int probe_driver_stop_called = 0;

static error_t probe_driver_probe(Device*) { return probe_present ? ERROR_NONE : ERROR_NOT_FOUND; }
static error_t probe_driver_start(Device*) { probe_driver_start_called++; return ERROR_NONE; }
static error_t probe_driver_stop(Device*) { probe_driver_stop_called++; return ERROR_NONE; }

static Driver probe_driver = {
    .name = "hotplug_test_probe_driver",
    .compatible = (const char*[]) { "hotplug_test,probe", nullptr },
    .start_device = probe_driver_start,
    .stop_device = probe_driver_stop,
    .probe = probe_driver_probe,
    .api = nullptr,
    .device_type = nullptr,
    .internal = nullptr,
};

TEST_CASE("device_hotplug_poll_once starts/stops a probe-capable device after debounce") {
    probe_present = false;
    probe_driver_start_called = 0;
    probe_driver_stop_called = 0;

    static Device device {
        .name = "hotplug_probe_device",
        .config = nullptr,
        .parent = nullptr,
    };

    CHECK_EQ(driver_construct_add(&probe_driver), ERROR_NONE);
    CHECK_EQ(device_construct_add(&device, "hotplug_test,probe"), ERROR_NONE);
    device_hotplug_register(&device);

    // Not present: repeated polling never starts it.
    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK_EQ(probe_driver_start_called, 0);
    CHECK_FALSE(device_is_ready(&device));

    // Present: needs DEVICE_HOTPLUG_DEBOUNCE_COUNT consecutive polls before it actually starts.
    probe_present = true;
    for (int i = 0; i < DEVICE_HOTPLUG_DEBOUNCE_COUNT - 1; i++) {
        CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
        CHECK_FALSE(device_is_ready(&device));
    }
    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK(device_is_ready(&device));
    CHECK_EQ(probe_driver_start_called, 1);

    // Already started: further polls while still present don't restart it.
    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK_EQ(probe_driver_start_called, 1);

    // Absent again: needs the same debounce before it stops.
    probe_present = false;
    for (int i = 0; i < DEVICE_HOTPLUG_DEBOUNCE_COUNT - 1; i++) {
        CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
        CHECK(device_is_ready(&device));
    }
    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK_FALSE(device_is_ready(&device));
    CHECK_EQ(probe_driver_stop_called, 1);

    // Only started/stopped, never removed/destructed.
    CHECK(device_is_added(&device));

    device_hotplug_unregister(&device);
    CHECK_EQ(device_remove(&device), ERROR_NONE);
    CHECK_EQ(device_destruct(&device), ERROR_NONE);
    CHECK_EQ(driver_remove_destruct(&probe_driver), ERROR_NONE);
}

TEST_CASE("device_start skips probe() for a device that isn't registered with the hotplug poller") {
    probe_present = false;
    probe_driver_start_called = 0;

    static Device device {
        .name = "hotplug_unregistered_device",
        .config = nullptr,
        .parent = nullptr,
    };

    CHECK_EQ(driver_construct_add(&probe_driver), ERROR_NONE);
    CHECK_EQ(device_construct_add(&device, "hotplug_test,probe"), ERROR_NONE);

    // Like a devicetree node with status "disabled" that is started explicitly
    CHECK_EQ(device_start(&device), ERROR_NONE);
    CHECK(device_is_ready(&device));
    CHECK_EQ(probe_driver_start_called, 1);

    // Registered devices are probed: an absent one doesn't start
    CHECK_EQ(device_stop(&device), ERROR_NONE);
    device_hotplug_register(&device);
    CHECK_NE(device_start(&device), ERROR_NONE);
    CHECK_FALSE(device_is_ready(&device));

    device_hotplug_unregister(&device);
    CHECK_EQ(device_remove(&device), ERROR_NONE);
    CHECK_EQ(device_destruct(&device), ERROR_NONE);
    CHECK_EQ(driver_remove_destruct(&probe_driver), ERROR_NONE);
}

static int mandatory_driver_start_called = 0;

static error_t mandatory_driver_start(Device*) { mandatory_driver_start_called++; return ERROR_NONE; }
static error_t mandatory_driver_stop(Device*) { return ERROR_NONE; }

static Driver mandatory_driver = {
    .name = "hotplug_test_mandatory_driver",
    .compatible = (const char*[]) { "hotplug_test,mandatory", nullptr },
    .start_device = mandatory_driver_start,
    .stop_device = mandatory_driver_stop,
    .probe = nullptr,
    .api = nullptr,
    .device_type = nullptr,
    .internal = nullptr,
};

TEST_CASE("device_hotplug_poll_once starts a mandatory (probe == NULL) device unconditionally on the first call") {
    mandatory_driver_start_called = 0;

    static Device device {
        .name = "hotplug_mandatory_device",
        .config = nullptr,
        .parent = nullptr,
    };

    CHECK_EQ(driver_construct_add(&mandatory_driver), ERROR_NONE);
    CHECK_EQ(device_construct_add(&device, "hotplug_test,mandatory"), ERROR_NONE);
    device_hotplug_register(&device);

    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK(device_is_ready(&device));
    CHECK_EQ(mandatory_driver_start_called, 1);

    // Already started: further polls don't restart it.
    CHECK_EQ(device_hotplug_poll_once(), ERROR_NONE);
    CHECK_EQ(mandatory_driver_start_called, 1);

    device_hotplug_unregister(&device);
    CHECK_EQ(device_stop(&device), ERROR_NONE);
    CHECK_EQ(device_remove(&device), ERROR_NONE);
    CHECK_EQ(device_destruct(&device), ERROR_NONE);
    CHECK_EQ(driver_remove_destruct(&mandatory_driver), ERROR_NONE);
}

static std::atomic<bool> blocking_probe_block { false };
static std::atomic<bool> blocking_probe_entered { false };
static std::atomic<bool> blocking_probe_release { false };
static std::atomic<int> blocking_driver_stop_called { 0 };
static std::atomic<bool> blocking_probe_saw_stop { false };

static error_t blocking_driver_probe(Device*) {
    if (blocking_probe_block) {
        blocking_probe_entered = true;
        while (!blocking_probe_release) {
            delay_millis(1);
        }
        blocking_probe_saw_stop = blocking_driver_stop_called > 0;
    }
    return ERROR_NONE;
}
static error_t blocking_driver_start(Device*) { return ERROR_NONE; }
static error_t blocking_driver_stop(Device*) { blocking_driver_stop_called++; return ERROR_NONE; }

static Driver blocking_probe_driver = {
    .name = "hotplug_test_blocking_probe_driver",
    .compatible = (const char*[]) { "hotplug_test,blocking_probe", nullptr },
    .start_device = blocking_driver_start,
    .stop_device = blocking_driver_stop,
    .probe = blocking_driver_probe,
    .api = nullptr,
    .device_type = nullptr,
    .internal = nullptr,
};

TEST_CASE("device_stop waits for a probe() that is in progress") {
    static Device device {
        .name = "hotplug_blocking_probe_device",
        .config = nullptr,
        .parent = nullptr,
    };
    blocking_probe_block = false;
    blocking_probe_entered = false;
    blocking_probe_release = false;
    blocking_driver_stop_called = 0;
    blocking_probe_saw_stop = false;

    CHECK_EQ(driver_construct_add(&blocking_probe_driver), ERROR_NONE);
    CHECK_EQ(device_construct_add(&device, "hotplug_test,blocking_probe"), ERROR_NONE);
    device_hotplug_register(&device);
    CHECK_EQ(device_start(&device), ERROR_NONE);

    blocking_probe_block = true;
    auto* poller = thread_alloc_full("hotplug_poller", 4096, [](void*) -> int32_t {
        device_hotplug_poll_once();
        return 0;
    }, nullptr, -1);
    CHECK_EQ(thread_start(poller), ERROR_NONE);
    while (!blocking_probe_entered) {
        delay_millis(1);
    }

    auto* stopper = thread_alloc_full("hotplug_stopper", 4096, [](void*) -> int32_t {
        return device_stop(&device);
    }, nullptr, -1);
    CHECK_EQ(thread_start(stopper), ERROR_NONE);
    delay_millis(20);
    CHECK_EQ(blocking_driver_stop_called.load(), 0);

    blocking_probe_release = true;
    CHECK_EQ(thread_join(poller, 1000, 1), ERROR_NONE);
    CHECK_EQ(thread_join(stopper, 1000, 1), ERROR_NONE);
    CHECK_EQ(thread_get_return_code(stopper), ERROR_NONE);
    thread_free(poller);
    thread_free(stopper);
    CHECK_FALSE(blocking_probe_saw_stop.load());
    CHECK_EQ(blocking_driver_stop_called.load(), 1);

    device_hotplug_unregister(&device);
    CHECK_EQ(device_remove(&device), ERROR_NONE);
    CHECK_EQ(device_destruct(&device), ERROR_NONE);
    CHECK_EQ(driver_remove_destruct(&blocking_probe_driver), ERROR_NONE);
}