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

#include <chrono>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <thread>

// ── FrameSourceFunc ───────────────────────────────────────────────────────────
// KPN source node: reads a movie file and emits one Frame per sample interval.
//
// Decode backend: FFmpeg hwaccel (CUDA/VAAPI, runtime-detected) when
// available, CPU otherwise.
//
// Sampling strategy: seek to the next target timestamp rather than decoding
// every frame, which is fast even for 1-FPS sampling of a 2-hour film.
//
// EOF handling: when the movie ends, emits a Frame with eof=true, then sleeps
// 500 ms between subsequent calls until the KPN network stops the thread.

struct FrameSourceFunc {
    static constexpr std::string_view label() { return "frame_source"; }

    explicit FrameSourceFunc(const Config& cfg)
        : decoder_(std::make_unique<FFmpegDecoder>(
              cfg.movie_path, /*use_hw=*/true,
              /*out_scale=*/cfg.scene_detect ? cfg.dense_scale : 1.0f))
    {
        sample_interval_sec_ = 1.0 / cfg.sample_fps;
        next_pos_sec_        = cfg.start_sec;
        end_sec_             = cfg.end_sec;
        max_decode_fps_      = cfg.max_decode_fps;

        // Dense mode: emit every native-rate frame instead of seeking to each
        // 1-FPS sample point. Required by the TransNetV2 scene detector, which
        // needs consecutive frames. A downstream decimator drops back to
        // sample_fps for the face pipeline. When dense, we advance by the
        // decoder's frame period (best effort — read_at decodes forward past the
        // last position, so consecutive small steps yield consecutive frames).
        dense_ = cfg.scene_detect;
        double dense_fps = 0.0;
        if (dense_) {
            double vfps = decoder_->fps();
            if (vfps <= 0.0) vfps = 25.0;
            // Dense decode rate: capped at native fps. A lower scene_decode_fps
            // decodes fewer frames (big speedup); TransNetV2 still localises cuts
            // and boundary timestamps stay exact (keyed off real timestamps).
            dense_fps = (cfg.scene_decode_fps > 0.f)
                      ? std::min(static_cast<double>(cfg.scene_decode_fps), vfps)
                      : vfps;
            dense_step_sec_ = 1.0 / dense_fps;
            if (cfg.dense_scale > 0.f && cfg.dense_scale < 1.f)
                bbox_upscale_ = 1.0f / cfg.dense_scale;
        }

        double total_s = decoder_->duration_sec();
        double span_s  = (end_sec_ > 0 ? std::min(end_sec_, total_s) : total_s)
                         - cfg.start_sec;
        double emit_fps = dense_ ? dense_fps : cfg.sample_fps;
        int n_frames = static_cast<int>(span_s * emit_fps);
        std::cerr << "[frame_source] decoder=" << decoder_->codec_name()
                  << " (" << decoder_->hw_backend() << ")"
                  << "  video_fps=" << decoder_->fps()
                  << (dense_ ? "  DENSE@" + std::to_string(dense_fps) + "fps" : "")
                  << "  start=" << cfg.start_sec << "s"
                  << (end_sec_ > 0 ? "  end=" + std::to_string(end_sec_) + "s" : "")
                  << "  sample_fps=" << cfg.sample_fps
                  << "  frames_to_emit=" << n_frames << "\n";
    }

    Frame operator()() {
        if (hit_eof_) {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
            return Frame{{}, 0.0, -1, /*eof=*/true};
        }

        // Wall-clock rate cap. KPN source nodes resubmit immediately on push
        // overflow, with no backpressure; without this cap we'd decode-and-drop
        // in a tight loop whenever downstream stalls. The cap also protects
        // ORT-only deployments where the pipeline can't keep up at decode speed.
        if (max_decode_fps_ > 0.f) {
            const auto now = std::chrono::steady_clock::now();
            if (!rate_started_) {
                rate_started_ = true;
                next_decode_at_ = now;
            }
            if (now < next_decode_at_)
                std::this_thread::sleep_until(next_decode_at_);
            const auto period = std::chrono::nanoseconds(
                static_cast<int64_t>(1e9f / max_decode_fps_));
            // Anchor the next slot off the slot we just consumed, not off
            // wall-clock now() — keeps the average rate stable. If we fell
            // behind by more than one period, snap forward to avoid building
            // up an unbounded sleep debt.
            next_decode_at_ += period;
            if (next_decode_at_ < now)
                next_decode_at_ = now + period;
        }

        auto t0 = std::chrono::steady_clock::now();
        cv::Mat img = decoder_->read_at(next_pos_sec_);
        auto t1 = std::chrono::steady_clock::now();
        double decode_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
        decode_ms_acc_  += decode_ms;
        ++decode_count_;
        if (decode_count_ % 10 == 0) {
            double avg_ms  = decode_ms_acc_ / 10.0;
            double avg_fps = avg_ms > 0.0 ? 1000.0 / avg_ms : 0.0;
            std::cerr << "[frame_source] decode  avg=" << avg_ms << "ms"
                      << "  fps=" << avg_fps << "\n";
            decode_ms_acc_ = 0.0;
        }

        if (img.empty()) {
            hit_eof_ = true;
            std::cerr << "[frame_source] EOF at t=" << next_pos_sec_ << "s\n";
            return Frame{{}, next_pos_sec_, frame_idx_++, /*eof=*/true};
        }

        // Cut detection is a downstream concern: camera_position_change_detector
        // owns the histogram compare and sets Frame::is_cut. The source emits
        // is_cut=false and only decodes/samples frames.
        Frame f{img, next_pos_sec_, frame_idx_++, /*eof=*/false, /*is_cut=*/false};
        // When dense_scale downscaled the decode, tell downstream how to map
        // detector coordinates back to original video resolution.
        f.bbox_upscale = bbox_upscale_;
        next_pos_sec_ += dense_ ? dense_step_sec_ : sample_interval_sec_;
        if (end_sec_ > 0 && next_pos_sec_ > end_sec_) {
            hit_eof_ = true;
            std::cerr << "[frame_source] reached end_sec=" << end_sec_ << "s\n";
        }
        return f;
    }

private:
    std::unique_ptr<FFmpegDecoder> decoder_;
    double           sample_interval_sec_{1.0};
    bool             dense_{false};
    float            bbox_upscale_{1.f};
    double           dense_step_sec_{0.04};
    double           next_pos_sec_{0.0};
    double           end_sec_{-1.0};
    float            max_decode_fps_{0.f};
    std::chrono::steady_clock::time_point next_decode_at_{};
    bool             rate_started_{false};
    int64_t          frame_idx_{0};
    bool             hit_eof_{false};
    double           decode_ms_acc_{0.0};
    int              decode_count_{0};
};
