#include "fgc/ImagePipeline.h"
#include "fgc/mock/MockCameraSource.h"
#include "jxl_test_decode.h"

#include <doctest/doctest.h>

#include <chrono>
#include <filesystem>
#include <mutex>
#include <thread>

using namespace fgc;
using fgc_test::decode;
namespace fs = std::filesystem;

namespace {

struct FakeChannel : IControlChannel {
    std::mutex            m;
    std::vector<CamEvent> events;
    bool connect() override { return true; }
    void disconnect() override {}
    bool connected() const override { return true; }
    void publishStatus(int) override {}
    void publishCamEvent(const CamEvent& e) override {
        std::lock_guard<std::mutex> l(m);
        events.push_back(e);
    }
    size_t count() {
        std::lock_guard<std::mutex> l(m);
        return events.size();
    }
};

struct TempDir {
    fs::path path;
    TempDir() {
        path = fs::temp_directory_path() /
               ("fgc_pipe_" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()));
        fs::create_directories(path);
    }
    ~TempDir() { fs::remove_all(path); }
};

// A uniformly coloured raw frame: every red site `r`, green `g`, blue `b`
// (on the layout's own scale), in BayerRG order (R G / G B).
Frame bayerFrame(PixelLayout layout, uint32_t w, uint32_t h, int r, int g, int b, long long ts) {
    Frame f;
    f.width = w; f.height = h; f.layout = layout; f.channels = 1; f.timestamp_ms = ts;
    std::vector<uint16_t> m(static_cast<size_t>(w) * h);
    for (uint32_t y = 0; y < h; ++y)
        for (uint32_t x = 0; x < w; ++x)
            m[static_cast<size_t>(y) * w + x] = static_cast<uint16_t>(
                y % 2 == 0 ? (x % 2 == 0 ? r : g) : (x % 2 == 0 ? g : b));
    f.data.assign(layoutFrameBytes(layout, w, h), 0);
    if (layout == PixelLayout::BayerRG8) {
        for (size_t i = 0; i < m.size(); ++i) f.data[i] = static_cast<uint8_t>(m[i]);
    } else {
        for (size_t i = 0, o = 0; i < m.size(); i += 2, o += 3) {
            f.data[o]     = static_cast<uint8_t>(m[i] & 0xFF);
            f.data[o + 1] = static_cast<uint8_t>(((m[i] >> 8) & 0x0F) | ((m[i + 1] & 0x0F) << 4));
            f.data[o + 2] = static_cast<uint8_t>(m[i + 1] >> 4);
        }
    }
    return f;
}

bool waitSaved(const ImagePipeline& p, long long n, int ms = 20000) {
    const auto end = std::chrono::steady_clock::now() + std::chrono::milliseconds(ms);
    while (p.imagesSaved() < n && std::chrono::steady_clock::now() < end)
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
    return p.imagesSaved() >= n;
}

// Mean of channel c over the interior (de-Bayering blurs the frame's edge).
double interiorMean(const fgc_test::Decoded& d, int w, int h, int c) {
    double s = 0;
    size_t n = 0;
    for (int y = 4; y < h - 4; ++y)
        for (int x = 4; x < w - 4; ++x, ++n) s += d.px[(static_cast<size_t>(y) * w + x) * 3 + c];
    return s / n;
}

ImagePipeline::Params params(const TempDir& t, double distance = 0.0) {
    ImagePipeline::Params p;
    p.output_dir = t.path.string();
    p.tower      = "test";
    p.cameras    = {{"RGB", 0, distance, 1, true}};
    p.jxl_threads = 2;
    return p;
}

std::string onlyFile(const TempDir& t) {
    for (auto& e : fs::recursive_directory_iterator(t.path))
        if (e.path().extension() == ".jxl") return e.path().string();
    return {};
}

}  // namespace

TEST_CASE("a 12-bit raw frame is saved as a 12-bit colour file, colours in place") {
    TempDir     t;
    FakeChannel ch;
    ImagePipeline pipe(ch, [] { return Orientation{12.5f, 3.0f}; }, params(t));
    pipe.start();
    // Red 3200, green 1600, blue 800 of 4095.
    pipe.submit(bayerFrame(PixelLayout::BayerRG12p, 64, 48, 3200, 1600, 800, 1000));
    REQUIRE(waitSaved(pipe, 1));
    pipe.stop();
    const auto d = decode(onlyFile(t), 3);
    CHECK(d.bits == 12);
    REQUIRE(d.px.size() == 64u * 48 * 3);
    CHECK(interiorMean(d, 64, 48, 0) == doctest::Approx(3200).epsilon(0.01));
    CHECK(interiorMean(d, 64, 48, 1) == doctest::Approx(1600).epsilon(0.01));
    CHECK(interiorMean(d, 64, 48, 2) == doctest::Approx(800).epsilon(0.01));
    REQUIRE(ch.count() == 1);
    CHECK(ch.events[0].heading_decideg == 125);
    CHECK(ch.events[0].timestamp_ms == 1000);
}

TEST_CASE("an 8-bit raw frame is saved as 8-bit colour, either de-Bayer method") {
    for (bool ea : {true, false}) {
        CAPTURE(ea);
        TempDir     t;
        FakeChannel ch;
        auto        p = params(t);
        p.cameras[0].edge_aware = ea;
        ImagePipeline pipe(ch, nullptr, p);
        pipe.start();
        pipe.submit(bayerFrame(PixelLayout::BayerRG8, 64, 48, 200, 100, 50, 2000));
        REQUIRE(waitSaved(pipe, 1));
        pipe.stop();
        const auto d = decode(onlyFile(t), 3);
        CHECK(d.bits == 8);
        CHECK(interiorMean(d, 64, 48, 0) == doctest::Approx(200).epsilon(0.02));
        CHECK(interiorMean(d, 64, 48, 2) == doctest::Approx(50).epsilon(0.05));
    }
}

TEST_CASE("RGB8 frames are saved as before") {
    TempDir     t;
    FakeChannel ch;
    ImagePipeline pipe(ch, nullptr, params(t, 0.8));
    MockCameraSource cam(1, 64, 48, PixelLayout::RGB8);
    cam.setFrameCallback([&pipe](const Frame& f) { pipe.submit(f); });
    pipe.start();
    REQUIRE(cam.trigger());
    REQUIRE(waitSaved(pipe, 1));
    pipe.stop();
    const auto d = decode(onlyFile(t), 3);
    CHECK(d.bits == 8);
    CHECK(d.px.size() == 64u * 48 * 3);
}

TEST_CASE("rotation applies to raw frames too") {
    TempDir     t;
    FakeChannel ch;
    auto        p = params(t);
    p.cameras[0].rotation_deg = 90;
    ImagePipeline pipe(ch, nullptr, p);
    pipe.start();
    pipe.submit(bayerFrame(PixelLayout::BayerRG12p, 64, 48, 3200, 1600, 800, 3000));
    REQUIRE(waitSaved(pipe, 1));
    pipe.stop();
    const auto d = decode(onlyFile(t), 3);
    CHECK(d.px.size() == 48u * 64 * 3);   // 48 wide, 64 high now
    CHECK(interiorMean(d, 48, 64, 0) == doctest::Approx(3200).epsilon(0.01));
}

TEST_CASE("several workers save every frame") {
    TempDir     t;
    FakeChannel ch;
    auto        p = params(t, 1.0);
    p.workers     = 3;
    p.jxl_threads = 1;
    ImagePipeline pipe(ch, nullptr, p);
    pipe.start();
    for (int i = 0; i < 9; ++i)
        pipe.submit(bayerFrame(PixelLayout::BayerRG12p, 64, 48, 3000, 1500, 700, 5000 + i));
    REQUIRE(waitSaved(pipe, 9));
    pipe.stop();
    CHECK(ch.count() == 9);
    size_t files = 0;
    for (auto& e : fs::recursive_directory_iterator(t.path)) files += e.path().extension() == ".jxl";
    CHECK(files == 9);
}

TEST_CASE("save timing counts the pictures saved in the last minute and how long they took") {
    TempDir     t;
    FakeChannel ch;
    ImagePipeline pipe(ch, nullptr, params(t, 1.0));
    CHECK(pipe.saveTiming().count == 0);
    pipe.start();
    for (int i = 0; i < 3; ++i)
        pipe.submit(bayerFrame(PixelLayout::BayerRG12p, 64, 48, 3000, 1500, 700, 7000 + i));
    REQUIRE(waitSaved(pipe, 3));
    pipe.stop();
    const auto s = pipe.saveTiming();
    CHECK(s.count == 3);
    CHECK(s.mean_ms > 0.0);
    CHECK(s.max_ms >= s.mean_ms);
}

TEST_CASE("canAccept says no once queue_max frames are waiting, yes once they are saved") {
    TempDir     t;
    FakeChannel ch;
    auto        p = params(t, 1.0);
    p.queue_max   = 2;
    ImagePipeline pipe(ch, nullptr, p);
    CHECK(pipe.canAccept());
    pipe.submit(bayerFrame(PixelLayout::BayerRG8, 32, 32, 100, 100, 100, 1));
    CHECK(pipe.canAccept());
    pipe.submit(bayerFrame(PixelLayout::BayerRG8, 32, 32, 100, 100, 100, 2));
    CHECK_FALSE(pipe.canAccept());       // not started: both still waiting
    CHECK(pipe.pending() == 2);
    pipe.submit(bayerFrame(PixelLayout::BayerRG8, 32, 32, 100, 100, 100, 3));  // never refused
    CHECK(pipe.pending() == 3);
    pipe.start();
    REQUIRE(waitSaved(pipe, 3));
    pipe.stop();
    CHECK(pipe.canAccept());
    CHECK(pipe.pending() == 0);
}

TEST_CASE("a frame shorter than its layout is dropped, not saved") {
    TempDir     t;
    FakeChannel ch;
    ImagePipeline pipe(ch, nullptr, params(t));
    pipe.start();
    Frame f = bayerFrame(PixelLayout::BayerRG12p, 64, 48, 1, 1, 1, 7);
    f.data.resize(f.data.size() / 2);
    pipe.submit(f);
    pipe.submit(bayerFrame(PixelLayout::BayerRG12p, 64, 48, 1, 1, 1, 8));
    REQUIRE(waitSaved(pipe, 1));
    pipe.stop();
    CHECK(pipe.imagesSaved() == 1);
    CHECK(ch.count() == 1);
}