#include <atomic>
#include <chrono>
#include <filesystem>
#include <fstream>
#include <stdexcept>
#include <thread>
#include <vector>
#include <cstring>

#include <destrum/Graphics/imageLoader.h>
#include <destrum/Serialization/SceneSerializer.h>
#include <destrum/Threading/JobSystem.h>

#include "TestUtils.h"

using namespace test;

namespace {
    JobSystem& Jobs() { return JobSystem::GetInstance(); }

    void testExecuteAndWait()
    {
        Jobs().Init(4);
        JobCounter counter;
        std::atomic<int> sum{0};
        for (int i = 1; i <= 100; ++i) {
            Jobs().Execute(counter, [&sum, i] { sum += i; });
        }
        Jobs().Wait(counter);
        check(counter.IsDone(), "counter should be done after Wait");
        check(sum == 5050, "every executed job should run before Wait returns");
        Jobs().Shutdown();
    }

    void testDispatchCoversEveryIndexOnce()
    {
        Jobs().Init(4);
        constexpr std::uint32_t count = 1003;
        std::vector<std::atomic<int>> hits(count);
        JobCounter counter;
        Jobs().Dispatch(counter, count, 64, [&hits](JobDispatchArgs args) {
            check(args.groupIndex == args.index / 64, "group index should match the index range");
            ++hits[args.index];
        });
        Jobs().Wait(counter);
        for (const auto& hit : hits) {
            check(hit == 1, "Dispatch should visit each index exactly once");
        }
        Jobs().Shutdown();
    }

    void testNestedWaitDoesNotDeadlock()
    {
        Jobs().Init(2);
        JobCounter outer;
        std::atomic<int> innerRuns{0};
        for (int i = 0; i < 8; ++i) {
            Jobs().Execute(outer, [&innerRuns] {
                JobCounter inner;
                for (int j = 0; j < 8; ++j) {
                    Jobs().Execute(inner, [&innerRuns] { ++innerRuns; });
                }
                Jobs().Wait(inner);
            });
        }
        Jobs().Wait(outer);
        check(innerRuns == 64, "nested jobs should all run");
        Jobs().Shutdown();
    }

    void testBlockedLowJobDoesNotStallHigh()
    {
        Jobs().Init(2);
        std::atomic<bool> release{false};
        auto blocker = Jobs().Async([&release] {
            while (!release) {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            }
            return 7;
        });

        JobCounter counter;
        std::atomic<int> runs{0};
        Jobs().Dispatch(counter, 32, 4, [&runs](JobDispatchArgs) { ++runs; });
        Jobs().Wait(counter);
        check(runs == 32, "high jobs should finish while a low job is blocked");

        release = true;
        check(blocker.get() == 7, "Async should return the job result");
        Jobs().Shutdown();
    }

    void testAsyncRethrows()
    {
        Jobs().Init(2);
        auto future = Jobs().Async([]() -> int { throw std::runtime_error("boom"); });
        bool threw = false;
        try {
            (void)future.get();
        } catch (const std::runtime_error&) {
            threw = true;
        }
        check(threw, "Async should rethrow the job's exception on get");

        JobCounter counter;
        Jobs().Execute(counter, [] { throw std::runtime_error("logged"); });
        Jobs().Wait(counter);
        check(counter.IsDone(), "a throwing job should still complete its counter");
        Jobs().Shutdown();
    }

    void testMainThreadQueue()
    {
        Jobs().Init(2);
        const auto mainId = std::this_thread::get_id();
        std::vector<int> order;
        bool onMain = true;

        JobCounter counter;
        Jobs().Execute(counter, [&] {
            for (int i = 0; i < 5; ++i) {
                Jobs().RunOnMainThread([&, i] {
                    order.push_back(i);
                    onMain = onMain && std::this_thread::get_id() == mainId;
                });
            }
        });
        Jobs().Wait(counter);
        check(order.empty(), "main thread jobs should wait for the drain");

        Jobs().DrainMainThreadQueue();
        check(order == std::vector<int>{0, 1, 2, 3, 4}, "main thread jobs should run in submission order");
        check(onMain, "main thread jobs should run on the draining thread");
        Jobs().Shutdown();
    }

    void testShutdownDrainsPendingJobs()
    {
        Jobs().Init(1);
        std::atomic<int> runs{0};
        std::vector<std::future<void>> futures;
        for (int i = 0; i < 50; ++i) {
            futures.push_back(Jobs().Async([&runs] { ++runs; }));
        }
        Jobs().Shutdown();
        Jobs().Shutdown();
        check(runs == 50, "Shutdown should finish queued jobs");
        check(!Jobs().IsRunning(), "job system should be stopped");
    }

    void testLoadSceneFileAsync()
    {
        Jobs().Init(2);
        const auto path = std::filesystem::temp_directory_path() / "destrum_async_scene.json";
        std::ofstream{path} << R"({"version":2,"name":"async","objects":[]})";

        auto good = SceneSerializer::LoadSceneFileAsync(path);
        auto missing = SceneSerializer::LoadSceneFileAsync(path.parent_path() / "destrum_no_such_scene.json");
        check(good.get().success, "async scene load should parse a valid file");
        check(!missing.get().success, "async scene load should report a missing file");

        std::filesystem::remove(path);
        Jobs().Shutdown();
    }

    void testAsyncHighPriority()
    {
        Jobs().Init(2);
        auto future = Jobs().Async([] { return std::this_thread::get_id(); }, JobPriority::High);
        check(future.get() != std::this_thread::get_id(), "high priority Async should run on a worker");
        Jobs().Shutdown();
    }

    void testImageDecodeOnWorkersMatchesSerial()
    {
        const std::filesystem::path root = DESTRUM_ENGINE_ASSETS_SRC_DIR;
        const std::vector<std::filesystem::path> paths{
            root / "238882.png",
            root / "char.jpg",
            root / "flight-helmet" / "FlightHelmet_Materials_GlassPlasticMat_Normal.png",
            root / "does_not_exist.png",
        };

        Jobs().Init(4);
        std::vector<std::future<ImageData>> decoded;
        for (const auto& path : paths) {
            decoded.push_back(Jobs().Async([path] { return util::loadImage(path, TextureIntent::ColorSrgb); },
                                           JobPriority::High));
        }

        for (std::size_t i = 0; i < paths.size(); ++i) {
            const ImageData parallel = decoded[i].get();
            const ImageData serial = util::loadImage(paths[i], TextureIntent::ColorSrgb);
            check(parallel.vkFormat == serial.vkFormat && parallel.width == serial.width &&
                  parallel.height == serial.height && parallel.byteSize == serial.byteSize,
                  "decoding on a worker should give the same image layout");
            if (serial.pixels != nullptr) {
                check(parallel.pixels != nullptr &&
                      std::memcmp(parallel.pixels, serial.pixels, serial.byteSize) == 0,
                      "decoding on a worker should give the same pixels");
            }
        }
        check(decoded.size() == 4, "every image should be decoded");
        Jobs().Shutdown();
    }

    void testRunsInlineWhenStopped()
    {
        check(!Jobs().IsRunning(), "job system should start stopped");
        JobCounter counter;
        int value = 0;
        Jobs().Execute(counter, [&value] { value = 3; });
        check(value == 3 && counter.IsDone(), "jobs should run inline without workers");
        check(Jobs().Async([] { return 5; }).get() == 5, "Async should run inline without workers");
    }
}

void RunJobsTests()
{
    testRunsInlineWhenStopped();
    testExecuteAndWait();
    testDispatchCoversEveryIndexOnce();
    testNestedWaitDoesNotDeadlock();
    testBlockedLowJobDoesNotStallHigh();
    testAsyncRethrows();
    testMainThreadQueue();
    testShutdownDrainsPendingJobs();
    testLoadSceneFileAsync();
    testAsyncHighPriority();
    testImageDecodeOnWorkersMatchesSerial();
}