screen-recorder/obs-manager.cpp

#include "obs-manager.hpp"
#include "settings-manager.hpp"

#include <obs/obs.h>
#include <gdkmm/display.h>
#include <glibmm/ustring.h>
#include <gdk/gdkx.h>
#include <X11/X.h>
#include <ctime>
#include <list>
#include <filesystem>
#include <stdexcept>

using namespace std;
using namespace Glib;

static void obs_render(void *param, uint32_t cx, uint32_t cy)
{
    obs_render_main_texture();
}

OBSManager::OBSManager()
{
    mPlugins = {
        "obs-ffmpeg.so",
        "obs-outputs.so",
        "obs-x264.so",
        "linux-v4l2.so",
        "linux-capture.so",
    };
}

OBSManager::~OBSManager()
{
    Cleanup();
}

ustring OBSManager::GetVersion()
{
    return ustring(obs_get_version_string());
}

void OBSManager::Initialize()
{
    if (isInitialized)
        return;
    
    sources = list<obs_source_t*>();
    auto settings = new SettingsManager();
    LoadSettings(settings);
    
    if (!obs_startup("en-US", NULL, NULL))
        throw runtime_error("Failed to initialize OBS");
    
    loadPlugins();
    printTypes();
    
    obs_video_info v = {};
    v.graphics_module = "libobs-opengl.so.0";
    v.fps_num = 30000;
    v.fps_den = 1001;
    v.base_width = PreviewWidth;
    v.base_height = PreviewHeight;
    v.output_width = PreviewWidth;
    v.output_height = PreviewHeight;
    v.output_format = VIDEO_FORMAT_NV12;
    v.adapter = 0;
    v.gpu_conversion = true;
    v.colorspace = VIDEO_CS_601;
    v.range = VIDEO_RANGE_PARTIAL;
    v.scale_type = OBS_SCALE_BICUBIC;

    obs_audio_info a = {};
    a.samples_per_sec = 44100;
    a.speakers = SPEAKERS_STEREO;

    obs_reset_video(&v);
    obs_reset_audio(&a);
    
    isInitialized = true;
}

void OBSManager::SetPreviewWindow(XID wid, Display *wdisplay)
{
    auto settings = new SettingsManager();
    
    gs_init_data init = {};
    init.cx = PreviewWidth;
    init.cy = PreviewHeight;
    init.format = GS_BGRA;
    init.zsformat = GS_ZS_NONE;
    init.window.id = wid;
    init.window.display = (void*)wdisplay;
    
    display = obs_display_create(&init, 0);
    if (display == nullptr)
        throw runtime_error("Failed to create display");
    
    obs_display_add_draw_callback(display, obs_render, nullptr);
    
    obs_scene_t *scene = obs_scene_create("scene1");
    if (scene == NULL)
        throw runtime_error("Couldn't create scene\n");
    
    if (settings->GetBool(SETTINGS_KEY_SCREEN_ENABLED))
    {
        auto source = CreateScreenSource();
        obs_scene_add(scene, source);
        sources.push_back(source);
    }
    
    if (settings->GetBool(SETTINGS_KEY_WEBCAM_ENABLED))
    {
        vec2 scale;
        vec2_set(&scale, 0.5f, 0.5f);
        
        auto source = CreateWebcamSource();
        auto item = obs_scene_add(scene, source);
        obs_sceneitem_set_scale(item, &scale);
        sources.push_back(source);
    }
    
    if (settings->GetBool(SETTINGS_KEY_AUDIO_ENABLED))
    {
        auto source = CreateAudioSource();
        obs_scene_add(scene, source);
        sources.push_back(source);
    }
    
    obs_set_output_source(0, obs_scene_get_source(scene));
}

void OBSManager::StartRecording()
{
    if (isRecording)
        return;
    
    auto s = new SettingsManager();
    
    obs_encoder_t *venc = obs_video_encoder_create("obs_x264", "x264", nullptr, nullptr);
    obs_encoder_t *aenc = obs_audio_encoder_create("ffmpeg_aac", "aac", nullptr, 0, nullptr);
    
    obs_encoder_set_video(venc, obs_get_video());
    obs_encoder_set_audio(aenc, obs_get_audio());
    
    string path = s->GetWithDefault(SETTINGS_KEY_OUTPUT_DIR, filesystem::current_path());
    if (path.back() != '/')
        path += "/";
    
    string fileName = path + "recording_" + to_string(time(0)) + ".mp4";
    
    obs_data_t *settings = obs_data_create();
    obs_data_set_string(settings, "directory", path.c_str());
    obs_data_set_string(settings, "url", fileName.c_str());
    
    output = obs_output_create("ffmpeg_output", "ffmpeg_output", nullptr, nullptr);
    obs_output_set_video_encoder(output, venc);
    obs_output_set_audio_encoder(output, aenc, 0);
    obs_output_update(output, settings);
    obs_output_set_media(output, obs_get_video(), obs_get_audio());
    
    obs_data_release(settings);
    if (!obs_output_start(output))
        throw runtime_error("Failed to start recording");
    
    isRecording = true;
}

void OBSManager::StopRecording()
{
    if (isRecording && output != nullptr)
        obs_output_stop(output);
    
    isRecording = false;
}

obs_source_t *OBSManager::CreateScreenSource()
{
    obs_data_t *settings = obs_data_create();
    obs_source_t *source = obs_source_create("xshm_input", "Screen Source", settings, NULL);
    if (source == NULL)
        throw runtime_error("Couldn't create screen source");
    
    obs_data_release(settings);
    return source;
}

obs_source_t *OBSManager::CreateWebcamSource()
{
    obs_data_t *settings = obs_data_create();
    obs_data_set_string(settings, "device_id", mWebcamDeviceID.c_str());

    obs_source_t *source = obs_source_create("v4l2_input", "Webcam Source", settings, NULL);
    if (source == NULL)
        throw runtime_error("Couldn't create webcam source");
    
    obs_data_release(settings);
    return source;
}

obs_source_t *OBSManager::CreateAudioSource()
{
    obs_data_t *settings = obs_data_create();
    obs_source_t *source = obs_source_create("audio_line", "Audio Source", settings, NULL);
    if (source == NULL)
        throw runtime_error("Couldn't create screen source");
    
    obs_data_release(settings);
    return source;
}

void OBSManager::Cleanup()
{
    if (!isInitialized)
        return;
    
    StopRecording();
    
    for (auto source : sources)
    {
        obs_source_remove(source);
    }
    
    if (display != nullptr)
        obs_display_destroy(display);
    
    obs_shutdown();
    isInitialized = false;
}

void OBSManager::LoadSettings(SettingsManager *settings)
{
    mPluginDir = settings->GetWithDefault(SETTINGS_KEY_PLUGIN_DIR, SETTINGS_DEFAULT_PLUGIN_DIR);
    mOutputDir = settings->GetWithDefault(SETTINGS_KEY_OUTPUT_DIR, std::filesystem::current_path());
    mWebcamDeviceID = settings->Get(SETTINGS_KEY_VIDEO_DEVICE_ID);
}

void OBSManager::loadPlugin(ustring name)
{
    obs_module_t *module;
    
    ustring path;
    path.append(mPluginDir);
    path.append(name);
    
    int res = obs_open_module(&module, path.c_str(), nullptr);
    if (res != MODULE_SUCCESS)
        throw runtime_error("Failed to open plugin");
    
    obs_init_module(module);
}


void OBSManager::loadPlugins()
{
    for (ustring plugin : mPlugins)
    {
        loadPlugin(plugin);
    }
    
    obs_post_load_modules();
}

void OBSManager::printTypes()
{
    const char *t;
    for (int i = 0; obs_enum_input_types(i, &t); ++i)
    {
        cout << "input type: " << t << endl;
    }
    
    for (int i = 0; obs_enum_output_types(i, &t); ++i)
    {
        cout << "output type: " << t << endl;
    }
    
    for (int i = 0; obs_enum_encoder_types(i, &t); ++i)
    {
        cout << "encoder type: " << t << endl;
    }
}