module;

export module lsp.test.scheduler.async_executor;

import std;

import lsp.scheduler.async_executor;

using namespace std::chrono_literals;

namespace
{
    using lsp::scheduler::async_executor::AsyncExecutor;
    using lsp::scheduler::async_executor::TaskStatus;

    void Expect(bool condition, const std::string& message)
    {
        if (!condition)
            throw std::runtime_error(message);
    }

    class SchedulerTestSuite
    {
    public:
        struct Entry
        {
            std::string name;
            std::function<void()> fn;
        };

        void Add(const std::string& name, std::function<void()> fn)
        {
            entries_.push_back({ name, std::move(fn) });
        }

        int RunAll()
        {
            int failures = 0;
            for (const auto& entry : entries_)
            {
                try
                {
                    entry.fn();
                    std::cout << "[PASS] " << entry.name << std::endl;
                }
                catch (const std::exception& e)
                {
                    ++failures;
                    std::cout << "[FAIL] " << entry.name << " -> " << e.what() << std::endl;
                }
            }
            std::cout << "\nTotal: " << entries_.size() << ", Failures: " << failures << std::endl;
            return failures;
        }

    private:
        std::vector<Entry> entries_;
    };

    class BlockingCallbackCopy
    {
    public:
        BlockingCallbackCopy(std::latch& copy_started, std::latch& release_copy) :
            copy_started_(copy_started), release_copy_(release_copy)
        {
        }

        BlockingCallbackCopy(const BlockingCallbackCopy& other) :
            copy_started_(other.copy_started_), release_copy_(other.release_copy_)
        {
            copy_started_.count_down();
            release_copy_.wait();
        }

        BlockingCallbackCopy(BlockingCallbackCopy&&) = default;

        void operator()(const lsp::scheduler::async_executor::TaskResult&) const {}

    private:
        std::latch& copy_started_;
        std::latch& release_copy_;
    };
}

export namespace lsp::test::scheduler::async_executor
{
    int Run()
    {
        SchedulerTestSuite suite;

        suite.Add("Rejects operations on a default handle", [] {
            lsp::scheduler::async_executor::TaskHandle handle;
            Expect(!handle.Valid(), "default handle should be invalid");
            Expect(!handle.Cancel(), "invalid handle cancellation should fail");
            Expect(!handle.Wait().has_value(), "invalid handle wait should be empty");
            Expect(!handle.TryGetResult().has_value(), "invalid handle result should be empty");
        });

        suite.Add("Retains completed task result through handle", [] {
            AsyncExecutor executor{ 1 };
            auto handle = executor.Submit("fast", [](std::stop_token) {
                return std::optional<std::string>{ "done" };
            });

            auto result = handle.Wait();
            Expect(handle.Valid(), "completed handle should remain valid");
            Expect(result && result->status == TaskStatus::kCompleted, "task should complete");
            Expect(result->value == "done", "task value should remain available");
            Expect(handle.TryGetResult().has_value(), "completed result should be non-blocking");
            Expect(!handle.Cancel(), "completed task cancellation should fail");
        });

        suite.Add("Waits for callback completion", [] {
            AsyncExecutor executor{ 1 };
            std::latch callback_started{ 1 };
            std::latch release_callback{ 1 };
            auto handle = executor.Submit(
                "callback.wait",
                [](std::stop_token) { return std::optional<std::string>{ "done" }; },
                [&](const lsp::scheduler::async_executor::TaskResult&) {
                    callback_started.count_down();
                    release_callback.wait();
                });

            callback_started.wait();
            auto waiter = std::async(std::launch::async, [&handle] { return handle.Wait(); });
            const bool blocked = waiter.wait_for(20ms) == std::future_status::timeout;
            const bool unavailable = !handle.TryGetResult().has_value();
            release_callback.count_down();
            auto result = waiter.get();

            Expect(blocked, "wait should block until callback returns");
            Expect(unavailable, "result should remain unavailable during callback");
            Expect(result && result->status == TaskStatus::kCompleted, "wait should return task result");
        });

        suite.Add("Skips a queued task cancelled through its handle", [] {
            AsyncExecutor executor{ 1 };
            std::latch blocker_started{ 1 };
            std::latch release_blocker{ 1 };
            auto blocker = executor.Submit("blocker", [&](std::stop_token) {
                blocker_started.count_down();
                release_blocker.wait();
                return std::optional<std::string>{};
            });

            blocker_started.wait();
            std::atomic<bool> ran{ false };
            auto queued = executor.Submit("queued", [&](std::stop_token) {
                ran = true;
                return std::optional<std::string>{};
            });
            const bool first_cancel = queued.Cancel();
            const bool second_cancel = queued.Cancel();
            release_blocker.count_down();

            auto queued_result = queued.Wait();
            auto blocker_result = blocker.Wait();
            Expect(first_cancel, "first queued cancellation should succeed");
            Expect(!second_cancel, "repeated cancellation should fail");
            Expect(queued_result && queued_result->status == TaskStatus::kCancelled, "queued task should cancel");
            Expect(!ran, "cancelled queued closure must not execute");
            Expect(blocker_result && blocker_result->status == TaskStatus::kCompleted, "blocker should finish");
        });

        suite.Add("Running task observes stop token", [] {
            AsyncExecutor executor{ 1 };
            std::latch started{ 1 };
            std::atomic<bool> observed_stop{ false };
            std::atomic<bool> release_task{ false };
            auto handle = executor.Submit("running", [&](std::stop_token stop_token) {
                started.count_down();
                while (!stop_token.stop_requested() && !release_task.load())
                    std::this_thread::yield();
                observed_stop = stop_token.stop_requested();
                return std::optional<std::string>{ "ignored" };
            });

            started.wait();
            const bool cancelled = executor.Cancel("running");
            release_task = true;
            auto result = handle.Wait();
            Expect(cancelled, "ID cancellation should find current task");
            Expect(observed_stop, "running task should observe stop token");
            Expect(result && result->status == TaskStatus::kCancelled, "running task should finish cancelled");
            Expect(!result->value.has_value(), "cancelled task should not expose a value");
        });

        suite.Add("Does not hold state lock while requesting stop", [] {
            AsyncExecutor executor{ 1 };
            lsp::scheduler::async_executor::TaskHandle handle;
            std::latch allow_callback_registration{ 1 };
            std::latch callback_registered{ 1 };
            std::atomic<bool> reentrant_cancel_blocked{ false };
            std::thread reentrant_cancel;

            handle = executor.Submit("stop.callback", [&](std::stop_token stop_token) {
                allow_callback_registration.wait();
                std::stop_callback callback(stop_token, [&] {
                    auto reentrant_done = std::make_shared<std::promise<void>>();
                    auto reentrant_future = reentrant_done->get_future();
                    reentrant_cancel = std::thread([&, reentrant_done] {
                        handle.Cancel();
                        reentrant_done->set_value();
                    });
                    reentrant_cancel_blocked = reentrant_future.wait_for(20ms) == std::future_status::timeout;
                });
                callback_registered.count_down();
                while (!stop_token.stop_requested())
                    std::this_thread::yield();
                return std::optional<std::string>{};
            });

            allow_callback_registration.count_down();
            callback_registered.wait();
            const bool cancelled = handle.Cancel();
            reentrant_cancel.join();
            auto result = handle.Wait();

            Expect(cancelled, "initial cancellation should succeed");
            Expect(!reentrant_cancel_blocked, "stop callback should reenter cancellation without blocking on state lock");
            Expect(result && result->status == TaskStatus::kCancelled, "task should finish cancelled");
        });

        suite.Add("Old same-ID handle cannot cancel replacement", [] {
            AsyncExecutor executor{ 2 };
            std::latch first_started{ 1 };
            std::atomic<bool> release_first{ false };
            auto first = executor.Submit("duplicate", [&](std::stop_token stop_token) {
                first_started.count_down();
                while (!stop_token.stop_requested() && !release_first.load())
                    std::this_thread::yield();
                return std::optional<std::string>{ "first" };
            });
            first_started.wait();

            std::latch second_started{ 1 };
            std::latch release_second{ 1 };
            std::atomic<bool> second_saw_stop{ false };
            auto second = executor.Submit("duplicate", [&](std::stop_token stop_token) {
                second_started.count_down();
                release_second.wait();
                second_saw_stop = stop_token.stop_requested();
                return std::optional<std::string>{ "second" };
            });
            second_started.wait();

            const bool old_cancel = first.Cancel();
            release_first = true;
            release_second.count_down();
            auto first_result = first.Wait();
            auto second_result = second.Wait();

            Expect(!old_cancel, "already replaced task should reject repeated cancellation");
            Expect(first_result && first_result->status == TaskStatus::kCancelled, "replaced task should cancel");
            Expect(second_result && second_result->status == TaskStatus::kCompleted, "replacement should complete");
            Expect(!second_saw_stop, "old handle must not stop replacement");
        });

        suite.Add("Counts every active same-ID instance", [] {
            AsyncExecutor executor{ 2 };
            std::latch first_started{ 1 };
            std::latch second_started{ 1 };
            std::latch release_tasks{ 1 };
            auto first = executor.Submit("same", [&](std::stop_token) {
                first_started.count_down();
                release_tasks.wait();
                return std::optional<std::string>{};
            });
            first_started.wait();
            auto second = executor.Submit("same", [&](std::stop_token) {
                second_started.count_down();
                release_tasks.wait();
                return std::optional<std::string>{};
            });
            second_started.wait();

            const auto active_count = executor.GetRunningTaskCount();
            release_tasks.count_down();
            executor.WaitAll();

            Expect(active_count == 2, "both same-ID instances should be active");
            Expect(executor.GetRunningTaskCount() == 0, "all instances should unregister");
            Expect(first.Wait().has_value() && second.Wait().has_value(), "both handles should retain results");
        });

        suite.Add("WaitAll observes a task before Taskflow submission", [] {
            AsyncExecutor executor{ 1 };
            std::latch callback_copy_started{ 1 };
            std::latch release_callback_copy{ 1 };
            AsyncExecutor::TaskCallback callback{
                BlockingCallbackCopy(callback_copy_started, release_callback_copy)
            };

            auto submitter = std::async(std::launch::async, [&] {
                return executor.Submit(
                    "submission.window",
                    [](std::stop_token) { return std::optional<std::string>{}; },
                    std::move(callback));
            });
            callback_copy_started.wait();
            auto waiter = std::async(std::launch::async, [&] { executor.WaitAll(); });
            const bool waited_for_registered_task = waiter.wait_for(20ms) == std::future_status::timeout;
            release_callback_copy.count_down();

            auto handle = submitter.get();
            waiter.get();
            auto result = handle.Wait();

            Expect(waited_for_registered_task, "WaitAll should not return while a registered task is being submitted");
            Expect(result && result->status == TaskStatus::kCompleted, "submitted task should complete");
        });

        suite.Add("Normalizes zero concurrency", [] {
            AsyncExecutor executor{ 0 };
            auto result = executor.Submit("zero", [](std::stop_token) {
                                      return std::optional<std::string>{ "done" };
                                  })
                              .Wait();
            Expect(result && result->status == TaskStatus::kCompleted, "zero concurrency should use one worker");
        });

        suite.Add("Preserves task exception", [] {
            AsyncExecutor executor{ 1 };
            auto result = executor.Submit("failure", [](std::stop_token) -> std::optional<std::string> {
                                      throw std::logic_error("original failure");
                                  })
                              .Wait();

            Expect(result && result->status == TaskStatus::kFailed, "throwing task should fail");
            Expect(result->error != nullptr, "failed task should retain exception");
            try
            {
                std::rethrow_exception(result->error);
            }
            catch (const std::logic_error& e)
            {
                Expect(std::string_view(e.what()) == "original failure", "original exception should be preserved");
            }
        });

        suite.Add("Ignores callback exception in task result", [] {
            AsyncExecutor executor{ 1 };
            auto handle = executor.Submit(
                "callback.failure",
                [](std::stop_token) { return std::optional<std::string>{ "done" }; },
                [](const lsp::scheduler::async_executor::TaskResult&) {
                    throw std::runtime_error("callback failure");
                });

            auto result = handle.Wait();
            Expect(result && result->status == TaskStatus::kCompleted, "callback failure should not alter task status");
            Expect(result->value == "done", "callback failure should not alter task value");
            Expect(result->error == nullptr, "callback failure should not become task error");
        });

        suite.Add("Reports terminal metrics", [] {
            AsyncExecutor executor{ 1 };
            std::latch blocker_started{ 1 };
            std::latch release_blocker{ 1 };
            auto completed = executor.Submit("complete", [&](std::stop_token) {
                blocker_started.count_down();
                release_blocker.wait();
                return std::optional<std::string>{ "done" };
            });
            blocker_started.wait();
            auto cancelled = executor.Submit("cancel", [](std::stop_token) {
                return std::optional<std::string>{};
            });
            auto failed = executor.Submit("fail", [](std::stop_token) -> std::optional<std::string> {
                throw std::runtime_error("failed");
            });
            const bool cancel_requested = cancelled.Cancel();
            release_blocker.count_down();
            executor.WaitAll();

            const auto metrics = executor.GetStatistics();
            Expect(cancel_requested, "queued task cancellation should succeed");
            Expect(completed.Wait()->status == TaskStatus::kCompleted, "completed status should match metrics");
            Expect(cancelled.Wait()->status == TaskStatus::kCancelled, "cancelled status should match metrics");
            Expect(failed.Wait()->status == TaskStatus::kFailed, "failed status should match metrics");
            Expect(metrics.running == 0, "no task should remain active");
            Expect(metrics.submitted == 3, "submitted count should be 3");
            Expect(metrics.completed == 1, "completed count should be 1");
            Expect(metrics.cancelled == 1, "cancelled count should be 1");
            Expect(metrics.failed == 1, "failed count should be 1");
        });

        const int failures = suite.RunAll();
        return failures == 0 ? 0 : 1;
    }
}