bytetrack-counter-cpp/src/video_writer.cpp

#include "video_writer.hpp"
#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <ctime>
#include <filesystem>
#include <sys/wait.h>
#include <unistd.h>

namespace fs = std::filesystem;

VideoBatchWriter::VideoBatchWriter(const std::string& shm_dir,
                                   const std::string& output_dir,
                                   int w, int h, int fps,
                                   float delay_sec,
                                   float end_delay_sec,
                                   int retention_days)
    : shm_dir_(shm_dir), output_dir_(output_dir),
      w_(w), h_(h), fps_(fps),
      delay_sec_(delay_sec), end_delay_sec_(end_delay_sec),
      retention_days_(retention_days)
{
    ensure_dir(shm_dir_);
    ensure_dir(output_dir_);

    for (const auto& entry : fs::directory_iterator(shm_dir_)) {
        if (entry.path().extension() == ".mp4") {
            fs::remove(entry.path());
        }
    }

    cleanup_retention();
}

VideoBatchWriter::~VideoBatchWriter() {
    shutdown();
}

void VideoBatchWriter::cleanup_retention() {
    auto cutoff = std::chrono::system_clock::now() - std::chrono::hours(retention_days_ * 24);
    auto cutoff_t = std::chrono::system_clock::to_time_t(cutoff);
    std::ostringstream cutoff_ss;
    cutoff_ss << std::put_time(std::localtime(&cutoff_t), "%Y-%m-%d");
    std::string cutoff_str = cutoff_ss.str();

    for (const auto& entry : fs::directory_iterator(output_dir_)) {
        if (!entry.is_directory()) continue;
        std::string dirname = entry.path().filename().string();
        if (dirname < cutoff_str) {
            std::error_code ec;
            fs::remove_all(entry.path(), ec);
        }
    }
}

void VideoBatchWriter::ensure_dir(const std::string& dir) {
    fs::create_directories(dir);
}

static std::string today_str() {
    auto t = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
    std::ostringstream ss;
    ss << std::put_time(std::localtime(&t), "%Y-%m-%d");
    return ss.str();
}

static std::string timestamp_str() {
    auto t = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
    std::ostringstream ss;
    ss << std::put_time(std::localtime(&t), "%H%M%S");
    return ss.str();
}

void VideoBatchWriter::note_motion() {
    if (state_ == IDLE) {
        start_recording();
        state_ = SPECULATING;
        state_start_ = std::chrono::steady_clock::now();
    }
}

void VideoBatchWriter::note_ayam_detection() {
    if (state_ == SPECULATING) {
        state_ = CONFIRMED;
    }
}

void VideoBatchWriter::on_batch_closed(int batch_number, const std::string& counting_date) {
    current_batch_ = batch_number;
    current_date_ = counting_date;
    if (state_ == CONFIRMED) {
        state_ = STOPPING;
        state_start_ = std::chrono::steady_clock::now();
    }
}

void VideoBatchWriter::feed_frame(const cv::Mat& raw) {
    auto now = std::chrono::steady_clock::now();

    if (state_ == SPECULATING) {
        float elapsed = std::chrono::duration<float>(now - state_start_).count();
        if (elapsed >= delay_sec_) {
            stop_and_delete();
            state_ = IDLE;
            return;
        }
    }

    if (state_ == STOPPING) {
        float elapsed = std::chrono::duration<float>(now - state_start_).count();
        if (elapsed >= end_delay_sec_) {
            stop_and_move(current_batch_);
            state_ = IDLE;
            return;
        }
    }

    if (state_ != IDLE && writer_.isOpened()) {
        writer_.write(raw);
    }
}

void VideoBatchWriter::start_recording() {
    if (writer_.isOpened()) writer_.release();

    auto t = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
    std::ostringstream name;
    name << "batch_pending_"
         << std::put_time(std::localtime(&t), "%Y%m%d_%H%M%S") << ".mp4";
    current_path_ = shm_dir_ + "/" + name.str();

    writer_.open(current_path_,
                 cv::VideoWriter::fourcc('a','v','c','1'),
                 fps_, cv::Size(w_, h_));

    if (writer_.isOpened()) {
        std::cout << "[RECORD] Speculating: " << current_path_ << std::endl;
    } else {
        std::cerr << "[RECORD] Failed to open: " << current_path_ << std::endl;
    }
}

static void copy_and_clean(const std::string& src, const std::string& dst) {
    if (src.empty() || dst.empty()) return;
    {
        std::ifstream ifs(src, std::ios::binary);
        if (ifs) {
            std::ofstream ofs(dst, std::ios::binary);
            ofs << ifs.rdbuf();
            ofs.flush();
        }
    }
    std::error_code ec;
    fs::remove(src, ec);
}

void VideoBatchWriter::stop_and_delete() {
    if (writer_.isOpened()) {
        writer_.release();
    }
    if (!current_path_.empty()) {
        std::cout << "[RECORD] Discarding: " << current_path_ << std::endl;
        std::error_code ec;
        fs::remove(current_path_, ec);
        current_path_.clear();
    }
}

void VideoBatchWriter::stop_and_move(int batch_number) {
    if (writer_.isOpened()) {
        writer_.release();
        std::cout << "[RECORD] Stopped: " << current_path_ << std::endl;
    }

    std::string src = current_path_;
    current_path_.clear();

    std::string day_dir = output_dir_ + "/" + (current_date_.empty() ? today_str() : current_date_);
    ensure_dir(day_dir);

    std::ostringstream name;
    name << "batch_" << batch_number << "_" << timestamp_str() << ".mp4";
    std::string dst = day_dir + "/" + name.str();

    pid_t pid = fork();
    if (pid == 0) {
        close(STDIN_FILENO);
        close(STDOUT_FILENO);
        close(STDERR_FILENO);
        copy_and_clean(src, dst);
        _exit(0);
    }
}

void VideoBatchWriter::shutdown() {
    if (writer_.isOpened()) {
        writer_.release();
    }

    auto prev_state = state_;
    state_ = IDLE;

    if (!current_path_.empty()) {
        std::string src = current_path_;
        current_path_.clear();

        if (prev_state == CONFIRMED || prev_state == STOPPING) {
            std::string day_dir = output_dir_ + "/" + (current_date_.empty() ? today_str() : current_date_);
            ensure_dir(day_dir);
            std::string dst = day_dir + "/" +
                fs::path(src).filename().string();

            pid_t pid = fork();
            if (pid == 0) {
                close(STDIN_FILENO);
                close(STDOUT_FILENO);
                close(STDERR_FILENO);
                copy_and_clean(src, dst);
                _exit(0);
            }
        } else {
            std::error_code ec;
            fs::remove(src, ec);
        }
    }
}