/*
 * Copyright 2024 Rive
 */

#include "common/test_harness.hpp"
#include "common/stacktrace.hpp"
#include "rive/rive_types.hpp"
#include "rive/math/math_types.hpp"
#include "tcp_client.hpp"
#include "png.h"
#include "zlib.h"
#include <vector>

#if defined(_WIN32) && !defined(NO_REDIRECT_OUTPUT)
#include <io.h>
static int pipe(int pipefd[2]) { return _pipe(pipefd, 65536, 0); }
#endif

#ifdef RIVE_ANDROID
#include <android/log.h>
#endif

constexpr static uint32_t REQUEST_TYPE_IMAGE_UPLOAD = 0;
constexpr static uint32_t REQUEST_TYPE_CLAIM_GM_TEST = 1;
constexpr static uint32_t REQUEST_TYPE_FETCH_RIV_FILE = 2;
constexpr static uint32_t REQUEST_TYPE_GET_INPUT = 3;
constexpr static uint32_t REQUEST_TYPE_CANCEL_INPUT = 4;
constexpr static uint32_t REQUEST_TYPE_PRINT_MESSAGE = 5;
constexpr static uint32_t REQUEST_TYPE_DISCONNECT = 6;
constexpr static uint32_t REQUEST_TYPE_APPLICATION_CRASH = 7;

static void check_early_exit()
{
    if (TestHarness::Instance().initialized())
    {
        // The tool should have called shutdown() before exit.
        TestHarness::Instance().onApplicationCrash("Early exit.");
    }
}

static void signal_wraper(const char* msg)
{
    TestHarness::Instance().onApplicationCrash(msg);
}

TestHarness& TestHarness::Instance()
{
    static TestHarness* instance = new TestHarness();
    return *instance;
}

TestHarness::TestHarness()
{
    stacktrace::replace_signal_handlers(signal_wraper, check_early_exit);
}

void TestHarness::init(std::unique_ptr<TCPClient> tcpClient,
                       size_t pngThreadCount)
{
    assert(!m_initialized);
    m_initialized = true;

    m_primaryTCPClient = std::move(tcpClient);

    initStdioThread();

    // We don't compile with emscripten pthreads.
#ifndef __EMSCRIPTEN__
    for (size_t i = 0; i < pngThreadCount; ++i)
    {
        m_encodeThreads.emplace_back(EncodePNGThread, this);
    }
#endif
}

void TestHarness::init(std::filesystem::path outputDir, size_t pngThreadCount)
{
    assert(!m_initialized);
    m_initialized = true;
    m_outputDir = outputDir;
#ifndef RIVE_REMOTE_ONLY
    std::filesystem::create_directories(m_outputDir);
#endif

#ifdef RIVE_ANDROID
    // Still pipe stdout and sterr to the android logs.
    initStdioThread();
#endif

    for (size_t i = 0; i < pngThreadCount; ++i)
    {
        m_encodeThreads.emplace_back(EncodePNGThread, this);
    }
}

void TestHarness::initStdioThread()
{
#if !defined(NO_REDIRECT_OUTPUT) && !defined(__EMSCRIPTEN__)
#ifndef _WIN32
    // Make stdout & stderr line buffered. (This is not supported on Windows.)
    setvbuf(stdout, NULL, _IOLBF, 0);
    setvbuf(stderr, NULL, _IOLBF, 0);
#endif
    // Pipe stdout and sterr back to the server.
    m_savedStdout = dup(1);
    m_savedStderr = dup(2);
    pipe(m_stdioPipe.data());
    dup2(m_stdioPipe[1], 1);
    dup2(m_stdioPipe[1], 2);
    m_stdioThread = std::thread(MonitorStdIOThread, this);
#endif
}

void TestHarness::monitorStdIOThread()
{
#if !defined(NO_REDIRECT_OUTPUT) && !defined(__EMSCRIPTEN__)
    assert(m_initialized);

    std::unique_ptr<TCPClient> threadTCPClient;
    if (m_primaryTCPClient != nullptr)
    {
        threadTCPClient = m_primaryTCPClient->clone();
    }

    char buff[1024];
    size_t readSize;
    while ((readSize = read(m_stdioPipe[0], buff, std::size(buff) - 1)) > 0)
    {
        buff[readSize] = '\0';
        if (threadTCPClient != nullptr)
        {
            threadTCPClient->send4(REQUEST_TYPE_PRINT_MESSAGE);
            threadTCPClient->sendString(buff);
        }
        else if (FILE* f = nullptr/*
                     fopen((m_outputDir / "rive_log.png").string().c_str(),
                           "a")*/)
        {
            // Sometimes it can help to also save a log file (e.g., when
            // ANDROID_LOG_DEBUG gets lost on browserstack).
            // DANGER: Writing this file may also cause failures, so turn the
            // if() back on with caution.
            fwrite(buff, 1, strlen(buff), f);
            fclose(f);
        }
#ifdef RIVE_ANDROID
        __android_log_write(ANDROID_LOG_DEBUG, "rive_android_tests", buff);
#endif
    }

    if (threadTCPClient != nullptr)
    {
        threadTCPClient->send4(REQUEST_TYPE_DISCONNECT);
        threadTCPClient->send4(false /* Don't shutdown the server yet */);
    }
#endif
}

void send_png_data_chunk(png_structp png, png_bytep data, png_size_t length)
{
    auto tcpClient = reinterpret_cast<TCPClient*>(png_get_io_ptr(png));
    tcpClient->send4(rive::math::lossless_numeric_cast<uint32_t>(length));
    tcpClient->sendall(data, length);
}

void flush_png_data(png_structp png) {}

static void save_png_impl(ImageSaveArgs args,
                          const std::filesystem::path& outputDir,
                          PNGCompression pngCompression,
                          TCPClient* tcpClient)
{
    assert(args.width > 0);
    assert(args.height > 0);
    std::string pngName = args.name + ".png";

    if (tcpClient == nullptr)
    {
        // We aren't connect to a test harness. Just save a file.
        auto destination = outputDir;
        destination /= pngName;
        destination.make_preferred();
        WritePNGFile(args.pixels.data(),
                     args.width,
                     args.height,
                     true,
                     destination.generic_string().c_str(),
                     pngCompression);
        return;
    }

    tcpClient->send4(REQUEST_TYPE_IMAGE_UPLOAD);
    tcpClient->sendString(pngName);

    auto png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (!png)
    {
        fprintf(stderr, "TestHarness: png_create_write_struct failed\n");
        abort();
    }

    // RLE with SUB gets best performance with our content.
    png_set_compression_level(png, 6);
    png_set_compression_strategy(png, Z_RLE);
    png_set_compression_strategy(png, Z_RLE);
    png_set_filter(png, 0, PNG_FILTER_SUB);

    auto info = png_create_info_struct(png);
    if (!info)
    {
        fprintf(stderr, "TestHarness: png_create_info_struct failed\n");
        abort();
    }

    if (setjmp(png_jmpbuf(png)))
    {
        fprintf(stderr, "TestHarness: Error during init_io\n");
        abort();
    }

    png_set_write_fn(png, tcpClient, &send_png_data_chunk, &flush_png_data);

    // Write header.
    if (setjmp(png_jmpbuf(png)))
    {
        fprintf(stderr, "TestHarness: Error during writing header\n");
        abort();
    }

    png_set_IHDR(png,
                 info,
                 args.width,
                 args.height,
                 8,
                 PNG_COLOR_TYPE_RGB_ALPHA,
                 PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE,
                 PNG_FILTER_TYPE_BASE);

    png_write_info(png, info);

    // Write bytes.
    if (setjmp(png_jmpbuf(png)))
    {
        fprintf(stderr, "TestHarness: Error during writing bytes\n");
        abort();
    }

    std::vector<uint8_t*> rows(args.height);
    for (uint32_t y = 0; y < args.height; ++y)
    {
        rows[y] = args.pixels.data() + (args.height - 1 - y) * args.width * 4;
    }
    png_write_image(png, rows.data());

    // End write.
    if (setjmp(png_jmpbuf(png)))
    {
        fprintf(stderr, "TestHarness: Error during end of write");
        abort();
    }

    png_write_end(png, NULL);
    png_destroy_write_struct(&png, &info);

    tcpClient->sendHandshake();
    tcpClient->recvHandshake();
}

void TestHarness::savePNG(ImageSaveArgs args)
{
    assert(m_initialized);
    if (!m_encodeThreads.empty())
    {
        m_encodeQueue.push(std::move(args));
    }
    else
    {
        save_png_impl(std::move(args),
                      m_outputDir,
                      m_pngCompression,
                      m_primaryTCPClient.get());
    }
}

void TestHarness::encodePNGThread()
{
    assert(m_initialized);

    std::unique_ptr<TCPClient> threadTCPClient;
    if (m_primaryTCPClient != nullptr)
    {
        threadTCPClient = m_primaryTCPClient->clone();
    }

    for (;;)
    {
        ImageSaveArgs args;
        m_encodeQueue.pop(args);
        if (args.quit)
        {
            break;
        }
        save_png_impl(std::move(args),
                      m_outputDir,
                      m_pngCompression,
                      threadTCPClient.get());
    }

    if (threadTCPClient != nullptr)
    {
        threadTCPClient->send4(REQUEST_TYPE_DISCONNECT);
        threadTCPClient->send4(false /* Don't shutdown the server yet */);
    }
}

bool TestHarness::claimGMTest(const std::string& name)
{
    if (m_primaryTCPClient != nullptr)
    {
        m_primaryTCPClient->send4(REQUEST_TYPE_CLAIM_GM_TEST);
        m_primaryTCPClient->sendString(name);
        return m_primaryTCPClient->recv4();
    }
    return true;
}

bool TestHarness::fetchRivFile(std::string& name, std::vector<uint8_t>& bytes)
{
    if (m_primaryTCPClient == nullptr)
    {
        return false;
    }
    m_primaryTCPClient->send4(REQUEST_TYPE_FETCH_RIV_FILE);
    uint32_t nameLength = m_primaryTCPClient->recv4();
    if (nameLength == TCPClient::SHUTDOWN_TOKEN)
    {
        return false;
    }
    name.resize(nameLength);
    m_primaryTCPClient->recvall(name.data(), nameLength);
    uint32_t fileSize = m_primaryTCPClient->recv4();
    bytes.resize(fileSize);
    m_primaryTCPClient->recvall(bytes.data(), fileSize);
    return true;
}

void TestHarness::printMessageOnServer(const char* msg)
{
    if (m_primaryTCPClient != nullptr)
    {
        m_primaryTCPClient->send4(REQUEST_TYPE_PRINT_MESSAGE);
        m_primaryTCPClient->sendString(msg);
    }
}

void TestHarness::inputPumpThread()
{
    assert(m_initialized);

    if (m_primaryTCPClient == nullptr)
    {
        return;
    }

    std::unique_ptr<TCPClient> threadTCPClient = m_primaryTCPClient->clone();

    for (std::vector<char> keys; m_initialized;)
    {
        threadTCPClient->send4(REQUEST_TYPE_GET_INPUT);
        uint32_t len = threadTCPClient->recv4();
        if (len == TCPClient::SHUTDOWN_TOKEN)
        {
            return;
        }
        keys.resize(len);
        threadTCPClient->recvall(keys.data(), len);
        for (char key : keys)
        {
            m_inputQueue.push(char(key));
        }
    }

    threadTCPClient->send4(REQUEST_TYPE_DISCONNECT);
    threadTCPClient->send4(false /* Don't shutdown the server yet */);
}

bool TestHarness::peekChar(char& key)
{
#ifdef __EMSCRIPTEN__
    // We don't compile with emscripten pthreads.
    return false;
#endif
    if (m_primaryTCPClient == nullptr)
    {
        return false;
    }
    if (!m_inputPumpThread.joinable())
    {
        m_inputPumpThread = std::thread(InputPumpThread, this);
    }
    return m_inputQueue.try_pop(key);
}

void TestHarness::shutdown()
{
    if (!m_initialized)
    {
        return;
    }

    for (std::thread& thread RIVE_MAYBE_UNUSED : m_encodeThreads)
    {
        m_encodeQueue.push({.quit = true});
    }
    for (std::thread& thread : m_encodeThreads)
    {
        thread.join();
    }
    m_encodeThreads.clear();

    shutdownStdioThread();
    shutdownInputPumpThread();

    if (m_primaryTCPClient != nullptr)
    {
        m_primaryTCPClient->send4(REQUEST_TYPE_DISCONNECT);
        m_primaryTCPClient->send4(true /* Shutdown the server */);
        m_primaryTCPClient = nullptr;
    }

    m_initialized = false;
}

void TestHarness::shutdownStdioThread()
{
#if !defined(NO_REDIRECT_OUTPUT) && !defined(__EMSCRIPTEN__)
    if (m_savedStdout != 0 || m_savedStderr != 0)
    {
        // Restore stdout and stderr.
        dup2(m_savedStdout, 1);
        dup2(m_savedStderr, 2);
        close(m_savedStdout);
        close(m_savedStderr);
        m_savedStdout = m_savedStderr = 0;

        // Shutdown the stdio-monitoring thread and pipe.
        close(m_stdioPipe[1]);
        m_stdioThread.join();
        close(m_stdioPipe[0]);
        m_stdioPipe = {0, 0};
    }
#endif
}

void TestHarness::shutdownInputPumpThread()
{
    if (m_inputPumpThread.joinable())
    {
        m_primaryTCPClient->send4(REQUEST_TYPE_CANCEL_INPUT);
        m_inputPumpThread.join();
    }
}

void TestHarness::onApplicationCrash(const char* message)
{
    if (m_primaryTCPClient != nullptr)
    {
        // Buy monitorStdIOThread() some time to finish pumping any messages
        // related to this abort.

        // std::this_thread::sleep_for causes weird link issues in unreal. just
        // use sleep instead
#if defined(RIVE_UNREAL) && defined(_WIN32)
        Sleep(100);
#else
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
#endif
        shutdownStdioThread();
        shutdownInputPumpThread();
        m_primaryTCPClient->send4(REQUEST_TYPE_APPLICATION_CRASH);
        m_primaryTCPClient->sendString(message);
    }
}