#pragma once
#include "types.hpp"
#include "config.hpp"
#include "scene_boundaries.hpp"

#include <memory>
#include "inference/scene_detector.hpp"

#include <opencv2/imgproc.hpp>   // cv::resize, for to_model_input

#include <nlohmann/json.hpp>
#include <algorithm>
#include <atomic>
#include <deque>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

// ── SceneDetectorFunc ─────────────────────────────────────────────────────────
// KPN sink node: TransNetV2 shot-boundary detection on the dense frame stream.
//
// Buffers incoming (dense, native-rate) Frames into a rolling window of
// ISceneDetector::kWindow (=100) frames. Every `stride` frames it runs one
// inference and reads back per-frame boundary probabilities, but only trusts the
// central region of each window — TransNetV2 (like most sliding-window boundary
// models) is unreliable near the window edges where it lacks temporal context.
// Overlapping windows by (kWindow - stride) frames means every frame is scored
// from at least one window's trusted centre.
//
// Boundaries (prob > scene_threshold, local maxima) are collected with their
// timestamps and written to scenes.json alongside the main annotations output on
// EOF. This branch is terminal: it produces no pipeline messages, only a file.

struct SceneDetectorFunc {
    static constexpr std::string_view label() { return "scene_detector"; }

    /// TRACES: AR-010 | SR-002
    /// Publish each window's verdict as it is scored, so the face branch — held
    /// back by channel depth — can consult it for frames it has not reached yet.
    void set_boundaries(std::shared_ptr<SceneBoundaries> b) { shared_ = std::move(b); }

    SceneDetectorFunc(const Config& cfg, std::atomic<bool>& done)
        : detector_(make_scene_detector(cfg))
        , threshold_(cfg.scene_threshold)
        , stride_(std::clamp(cfg.scene_stride, 1, ISceneDetector::kWindow))
        , output_path_(scenes_path(cfg.output_path))
        , movie_path_(cfg.movie_path)
        , done_(done)
    {
        // Trusted centre half of each window. Frames outside [guard, kWindow-guard)
        // are re-scored by an adjacent window, so we ignore them here to avoid
        // edge artefacts and double-counting.
        guard_ = (ISceneDetector::kWindow - stride_) / 2;
        std::cerr << "[scene_detector] threshold=" << threshold_
                  << "  stride=" << stride_
                  << "  guard=" << guard_
                  << "  output=" << output_path_ << "\n";
    }

    void operator()(Frame f) {
        if (f.eof) {
            flush_remaining();
            // Release anyone waiting on the join: the tail frames after the last
            // full window will never be covered, so waiting for them would hang.
            if (shared_) shared_->finish();
            write_output();
            done_.store(true, std::memory_order_release);
            return;
        }

        /// TRACES: AR-011 | SR-002
        // Learn the cadence of the stream from the stream itself, rather than
        // assuming one. See dedup_window_sec().
        if (prev_ts_ >= 0.0 && intervals_.size() < kCadenceSamples) {
            const double dt = f.timestamp_sec - prev_ts_;
            if (dt > 0.0) intervals_.push_back(dt);
        }
        prev_ts_ = f.timestamp_sec;

        images_.push_back(to_model_input(f.image));
        times_.push_back(f.timestamp_sec);

        // Once we have a full window, score it and slide forward by `stride`.
        while (static_cast<int>(images_.size()) >= ISceneDetector::kWindow) {
            score_window();
            for (int i = 0; i < stride_; ++i) {
                images_.pop_front();
                times_.pop_front();
            }
            window_base_ += stride_;
        }
    }

    /// TRACES: AR-004, AR-010 | SR-002
    /// Reduce a decoded frame to exactly what TransNetV2 consumes, once.
    ///
    /// The window used to hold the frames as decoded — full resolution — and
    /// leave the downscale to the backend. But the model's input is 48x27
    /// (`ISceneDetector::kFrameW/H`; the config note for `dense_scale` says so
    /// outright: "TransNetV2 downsamples to 48x27 regardless"), so the buffer
    /// held ~590 MB at 1080p to feed something that needs ~380 KB. That is not
    /// a channel capacity, so no amount of tuning channel depths would ever
    /// have found it.
    ///
    /// It is also redundant work. Windows overlap by `kWindow - stride`, so a
    /// frame appears in several of them and was re-downscaled once per window;
    /// now it is downscaled once, when it arrives.
    ///
    /// **This must reproduce the backends' preprocessing exactly**, because the
    /// project invariant is that every model gets the input it was trained for
    /// — a model run off-distribution returns confident, plausible, wrong
    /// output, and here that means fabricated shot boundaries. Both
    /// ort_backend.cpp and trt_backend.cpp guard mis-sized input with, in this
    /// order, `convertTo(CV_8UC3)` then
    /// `cv::resize(..., {kFrameW, kFrameH}, 0, 0, cv::INTER_AREA)`. The same
    /// two operations are done here, so the tensor the model receives is
    /// unchanged; the backend guard then sees a correctly-sized frame and does
    /// nothing. The interface has always specified this shape as the caller's
    /// job ("Each frame must already be kFrameW x kFrameH, BGR, CV_8UC3"), so
    /// this makes the node meet a contract it was already given.
    static cv::Mat to_model_input(const cv::Mat& src) {
        cv::Mat typed;
        if (src.type() != CV_8UC3) src.convertTo(typed, CV_8UC3);
        else                       typed = src;

        if (typed.cols == ISceneDetector::kFrameW &&
            typed.rows == ISceneDetector::kFrameH)
            return typed;

        cv::Mat small;
        cv::resize(typed, small, {ISceneDetector::kFrameW, ISceneDetector::kFrameH},
                   0, 0, cv::INTER_AREA);
        return small;
    }

    /// TRACES: AR-011 | SR-002
    // How close two boundaries have to be before they are the same boundary,
    // derived from the cadence the detector was actually fed.
    //
    // What this replaces is a literal 0.04 s — one frame at 25 fps, and silently
    // wrong at any other rate. On a 30 fps source it spans more than a frame, so
    // two cuts on consecutive frames merge into one and a real boundary is lost;
    // the output does not show this, it simply contains fewer cuts. Assuming a
    // frame rate is the same class of mistake as feeding a model the wrong rate,
    // which is why this belongs to AR-011 and not to a tidy-up.
    //
    // Half a frame, not a whole one, because the only thing being deduplicated is
    // one frame scored by two overlapping windows — a gap of zero. Two distinct
    // frames are a full interval apart and must both survive. Half an interval
    // separates those two cases without putting the decision on the knife-edge
    // where floating-point error settles it.
    //
    // Median, not mean: a seek, or a gap where the decoder dropped a frame,
    // contributes one long interval that would drag a mean and cannot move a
    // median.
    static double dedup_window_sec(std::vector<double> intervals) {
        if (intervals.empty()) return 0.0;   // <2 frames: nothing to deduplicate
        const std::size_t mid = intervals.size() / 2;
        std::nth_element(intervals.begin(), intervals.begin() + mid,
                         intervals.end());
        return intervals[mid] * 0.5;
    }

private:
    // Run TransNetV2 on the leading kWindow frames of the buffer and record any
    // boundaries found within the trusted centre region.
    void score_window() {
        std::vector<cv::Mat> win(images_.begin(),
                                 images_.begin() + ISceneDetector::kWindow);
        std::vector<float> probs = detector_->detect_window(win);

        // On the very first window there is no preceding window, so trust from 0;
        // otherwise skip the leading guard already covered by the previous window.
        const int lo = (window_base_ == 0) ? 0 : guard_;
        const int hi = ISceneDetector::kWindow - guard_;
        std::vector<double> fresh;
        for (int i = lo; i < hi; ++i) {
            if (probs[i] <= threshold_) continue;
            // Local maximum → the boundary frame (avoid a run of high scores
            // registering as several adjacent cuts).
            const bool peak =
                (i == 0        || probs[i] >= probs[i-1]) &&
                (i == kLast_() || probs[i] >= probs[i+1]);
            if (peak) {
                boundaries_.push_back({times_[i], probs[i]});
                fresh.push_back(times_[i]);
            }
        }

        /// TRACES: AR-010 | SR-002
        // Publish with a watermark: everything up to times_[hi-1] now has a
        // final verdict. The face branch consults this for frames it has not
        // reached yet, and the watermark is what lets it tell "no boundary
        // here" from "not scored yet".
        /// TRACES: AR-011 | SR-002
        // Hand the join the same dedup window scenes.json uses, derived from the
        // observed cadence rather than assumed. Set on every window because the
        // median refines as intervals accumulate; it converges within the first
        // window and costs a double assignment thereafter.
        if (shared_) {
            shared_->set_merge_window(dedup_window_sec(intervals_));
            if (hi > lo) shared_->publish(fresh, times_[hi - 1]);
        }
    }

    // At EOF the tail (< kWindow frames) never formed a full window. Pad it out
    // to kWindow by repeating the last frame so the final real frames still get
    // scored, then take only the region past what earlier windows covered.
    void flush_remaining() {
        const int n = static_cast<int>(images_.size());
        if (n == 0) return;
        std::vector<cv::Mat> win(images_.begin(), images_.end());
        cv::Mat last = win.back();
        while (static_cast<int>(win.size()) < ISceneDetector::kWindow)
            win.push_back(last);

        std::vector<float> probs = detector_->detect_window(win);
        const int lo = (window_base_ == 0) ? 0 : guard_;
        std::vector<double> fresh;
        for (int i = lo; i < n; ++i) {   // only real (non-padded) frames
            if (probs[i] <= threshold_) continue;
            const bool peak =
                (i == 0     || probs[i] >= probs[i-1]) &&
                (i == n - 1 || probs[i] >= probs[i+1]);
            if (peak) {
                boundaries_.push_back({times_[i], probs[i]});
                fresh.push_back(times_[i]);
            }
        }

        /// TRACES: AR-010 | SR-002
        // Publish the tail too. Without this the final frames — everything after
        // the last full window — reach the join with no verdict and are treated
        // as boundary-free without evidence, which is precisely the ambiguity
        // the watermark exists to prevent.
        if (shared_ && n > 0) {
            shared_->set_merge_window(dedup_window_sec(intervals_));
            shared_->publish(fresh, times_[n - 1]);
        }
    }

    void write_output() {
        if (written_) return;
        written_ = true;

        // Merge boundaries closer than one frame apart (dedup across window seams).
        const double dedup_sec = dedup_window_sec(intervals_);

        std::sort(boundaries_.begin(), boundaries_.end(),
                  [](const Boundary& a, const Boundary& b) {
                      return a.t < b.t;
                  });

        nlohmann::json root;
        root["schema_version"] = 1;
        root["movie"]          = movie_path_;
        root["model"]          = "transnetv2";
        root["threshold"]      = threshold_;
        nlohmann::json cuts = nlohmann::json::array();
        double last_t = -1e9;
        for (const auto& b : boundaries_) {
            if (b.t - last_t < dedup_sec) continue;
            cuts.push_back({{"t", b.t}, {"probability", b.prob}});
            last_t = b.t;
        }
        root["cuts"] = std::move(cuts);

        std::ofstream f(output_path_);
        if (!f.is_open()) {
            std::cerr << "\n[scene_detector] ERROR: cannot write "
                      << output_path_ << "\n";
            return;
        }
        f << root.dump(2) << "\n";
        // Report the derived cadence: VR-006 re-tunes scene_threshold against it,
        // and a rate that is not the source's is the first thing to suspect.
        std::cerr << "\n[scene_detector] wrote " << root["cuts"].size()
                  << " boundaries → " << output_path_
                  << "  (dedup=" << dedup_sec << "s from "
                  << (dedup_sec > 0.0 ? 0.5 / dedup_sec : 0.0) << " fps)\n";
    }

    static int kLast_() { return ISceneDetector::kWindow - 1; }

    // annotations.json → annotations.scenes.json (or scenes.json for bare names)
    static std::string scenes_path(const std::string& out) {
        auto dot = out.find_last_of('.');
        if (dot == std::string::npos) return out + ".scenes.json";
        return out.substr(0, dot) + ".scenes.json";
    }

    struct Boundary { double t; float prob; };

    // Enough to establish a rate; bounded so a feature-length film does not
    // accumulate one double per frame for a number that stops moving early.
    static constexpr std::size_t kCadenceSamples = 512;

    std::unique_ptr<ISceneDetector> detector_;
    float                 threshold_;
    int                   stride_;
    int                   guard_{0};
    std::string           output_path_;
    std::string           movie_path_;

    std::atomic<bool>&    done_;
    std::deque<cv::Mat>   images_;
    std::deque<double>    times_;
    int64_t               window_base_{0};   // frame index of images_.front()
    std::vector<Boundary> boundaries_;
    double                prev_ts_{-1.0};    // AR-011: cadence, learned not assumed
    std::vector<double>   intervals_;
    bool                  written_{false};
    std::shared_ptr<SceneBoundaries> shared_;   ///< AR-010 join point
};