#include "fgc/ImagePipeline.h"

#include "fgc/JpegXlEncoder.h"
#include "fgc/Logger.h"
#include "fgc/RawImage.h"

#include <algorithm>
#include <cerrno>
#include <chrono>
#include <filesystem>

#include <sys/resource.h>
#include <unistd.h>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>

namespace fs = std::filesystem;

namespace fgc {
namespace {

// Maps a clockwise rotation in degrees to the cv::rotate code. Returns false
// when there is nothing to do (0 deg, or an unexpected value - ConfigLoader
// rejects anything but 0/90/180/270 before it reaches here).
bool rotateCode(int deg, int& code) {
    switch (deg) {
        case 90:  code = cv::ROTATE_90_CLOCKWISE;        return true;
        case 180: code = cv::ROTATE_180;                 return true;
        case 270: code = cv::ROTATE_90_COUNTERCLOCKWISE; return true;
        default:  return false;
    }
}

}  // namespace

ImagePipeline::ImagePipeline(IControlChannel& channel, std::function<Orientation()> orientation,
                             Params params)
    : channel_(channel), orientation_(std::move(orientation)), params_(std::move(params)) {}

ImagePipeline::~ImagePipeline() { stop(); }

void ImagePipeline::start() {
    if (running_.exchange(true)) return;
    const int n = std::max(1, params_.workers);
    for (int i = 0; i < n; ++i) workers_.emplace_back(&ImagePipeline::run, this);
}

void ImagePipeline::stop() {
    if (!running_.exchange(false)) return;
    cv_.notify_all();
    for (auto& w : workers_)
        if (w.joinable()) w.join();
    workers_.clear();
}

bool ImagePipeline::canAccept() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return queue_.size() + in_flight_ < static_cast<size_t>(std::max(1, params_.queue_max));
}

size_t ImagePipeline::pending() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return queue_.size() + in_flight_;
}

void ImagePipeline::submit(const Frame& frame) {
    // Sample the pose now, while the frame is fresh, and carry it through the
    // queue. Both sources behind this callback are mutex-guarded snapshot reads,
    // so it costs the camera's delivery thread a lock and a copy.
    QueuedFrame q{frame, orientation_ ? orientation_() : Orientation{}};
    {
        std::lock_guard<std::mutex> lock(mutex_);
        queue_.push(std::move(q));
    }
    cv_.notify_one();
}

static long long steadyMs() {
    using namespace std::chrono;
    return duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
}

void ImagePipeline::noteSaveTime(double ms) {
    const long long now = steadyMs();
    std::lock_guard<std::mutex> lock(timing_mutex_);
    save_times_.emplace_back(now, ms);
    while (!save_times_.empty() && now - save_times_.front().first > 60000) save_times_.pop_front();
    if (timing_logged_ms_ == 0) timing_logged_ms_ = now;
    if (now - timing_logged_ms_ < 60000) return;
    timing_logged_ms_ = now;
    double sum = 0.0, max = 0.0;
    for (const auto& s : save_times_) {
        sum += s.second;
        max = std::max(max, s.second);
    }
    LOG_INFO << "saving: " << save_times_.size() << " picture"
             << (save_times_.size() == 1 ? "" : "s") << " in the last minute, "
             << static_cast<long>(sum / static_cast<double>(save_times_.size()) + 0.5)
             << " ms each on average (longest " << static_cast<long>(max + 0.5) << " ms)";
}

ImagePipeline::SaveTiming ImagePipeline::saveTiming() const {
    const long long now = steadyMs();
    std::lock_guard<std::mutex> lock(timing_mutex_);
    SaveTiming t;
    double sum = 0.0;
    for (const auto& s : save_times_) {
        if (now - s.first > 60000) continue;
        ++t.count;
        sum += s.second;
        t.max_ms = std::max(t.max_ms, s.second);
    }
    if (t.count > 0) t.mean_ms = sum / t.count;
    return t;
}

void ImagePipeline::run() {
    // Lower this worker's priority (Linux: nice is per thread, and the encoder's
    // own threads inherit it). Best effort: a failure only costs the head start.
    if (kWorkerNice != 0) {
        errno = 0;
        const int now = getpriority(PRIO_PROCESS, static_cast<id_t>(::gettid()));
        if (errno == 0 && now < kWorkerNice)
            setpriority(PRIO_PROCESS, static_cast<id_t>(::gettid()), kWorkerNice);
    }
    while (running_) {
        QueuedFrame queued;
        {
            std::unique_lock<std::mutex> lock(mutex_);
            cv_.wait(lock, [this] { return !queue_.empty() || !running_; });
            if (!running_ && queue_.empty()) break;
            queued = std::move(queue_.front());
            queue_.pop();
            ++in_flight_;
        }
        process(queued);
        std::lock_guard<std::mutex> lock(mutex_);
        --in_flight_;
    }
}

const ImagePipeline::CameraOutput& ImagePipeline::outputFor(int cam_id) const {
    if (cam_id >= 0 && cam_id < static_cast<int>(params_.cameras.size()))
        return params_.cameras[cam_id];
    // A source delivering an index we were not told about still has to write
    // somewhere identifiable rather than colliding with camera 0.
    static const CameraOutput fallback{"cam?", 0, 0.8, 4};
    return fallback;
}

void ImagePipeline::process(const QueuedFrame& queued) {
    const Frame& frame = queued.frame;
    const bool   known = frame.layout != PixelLayout::Unknown;
    if (!known && frame.channels != 1 && frame.channels != 3) {
        LOG_WARN << "Dropping frame with unsupported channel count " << frame.channels;
        return;
    }
    if (known && frame.data.size() < layoutFrameBytes(frame.layout, frame.width, frame.height)) {
        LOG_WARN << "Dropping " << layoutName(frame.layout) << " frame: " << frame.data.size()
                 << " bytes is short of " << frame.width << 'x' << frame.height;
        return;
    }

    const CameraOutput& out = outputFor(frame.cam_id);
    const std::string label = out.label.empty() ? ("cam" + std::to_string(frame.cam_id))
                                                : out.label;
    const fs::path     dir   = fs::path(params_.output_dir) / label;
    std::error_code    ec;
    fs::create_directories(dir, ec);
    if (ec) {
        LOG_ERROR << "Cannot create capture dir " << dir.string() << ": " << ec.message();
        return;
    }
    const fs::path file = dir / (std::to_string(frame.timestamp_ms) + ".jxl");
    const auto t0 = std::chrono::steady_clock::now();
    if (!save(frame, out, file.string())) return;
    noteSaveTime(std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0)
                     .count());

    // Sampled at submit(), not here: by this point the encode above has run and
    // the gimbal may already be slewing to the next waypoint.
    const Orientation& o = queued.orientation;
    CamEvent ev;
    ev.tower           = params_.tower;
    ev.camera          = label;
    ev.heading_decideg = static_cast<int>(o.yaw_deg * 10);
    ev.pitch_decideg   = static_cast<int>(o.pitch_deg * 10);
    ev.timestamp_ms    = frame.timestamp_ms;
    channel_.publishCamEvent(ev);
    {
        std::lock_guard<std::mutex> lock(last_event_mutex_);
        last_event_ = ev;
    }

    saved_.fetch_add(1);
    LOG_DEBUG << "Saved " << file.string() << " (" << frame.width << 'x' << frame.height << ' '
              << (known ? layoutName(frame.layout) : "legacy") << ')';
}

bool ImagePipeline::save(const Frame& frame, const CameraOutput& out, const std::string& file) {
    int  code   = 0;
    bool rotate = rotateCode(out.rotation_deg, code);
    const int h = static_cast<int>(frame.height), w = static_cast<int>(frame.width);
    auto* data  = const_cast<uint8_t*>(frame.data.data());

    cv::Mat img;   // what gets encoded: 8-bit (CV_8U) or 12-bit values in CV_16U
    int     bits = 8;
    if (frame.layout == PixelLayout::BayerRG8 || frame.layout == PixelLayout::BayerRG12p) {
        cv::Mat mosaic;
        if (frame.layout == PixelLayout::BayerRG8) {
            mosaic = cv::Mat(h, w, CV_8UC1, data);
        } else {
            mosaic.create(h, w, CV_16UC1);
            unpack12p(frame.data.data(), reinterpret_cast<uint16_t*>(mosaic.data),
                      static_cast<size_t>(w) * h);
            bits = 12;
        }
        // The camera's BayerRG is OpenCV's BayerBG (see RawImage.h).
        cv::cvtColor(mosaic, img,
                     out.edge_aware ? cv::COLOR_BayerBG2RGB_EA : cv::COLOR_BayerBG2RGB);
    } else {
        // RGB8, Mono8, or a legacy frame described by its channel count. The
        // Mat only wraps frame.data, so an unrotated frame costs no copy.
        img = cv::Mat(h, w, frame.channels == 3 ? CV_8UC3 : CV_8UC1, data);
    }
    if (rotate) {
        cv::Mat rotated;
        cv::rotate(img, rotated, code);
        img = rotated;
    }
    if (!img.isContinuous()) img = img.clone();

    if (!JpegXlEncoder::encodeToFile(file, img.data, img.cols, img.rows, img.channels(), bits,
                                     static_cast<float>(out.jxl_distance), out.jxl_effort,
                                     params_.jxl_threads)) {
        LOG_ERROR << "Failed to encode " << file;
        return false;
    }
    return true;
}

std::optional<CamEvent> ImagePipeline::lastEvent() const {
    std::lock_guard<std::mutex> lock(last_event_mutex_);
    return last_event_;
}

}  // namespace fgc