#include "gif-h/gif.h"
#include "Game.hpp"
#include "extension.hpp"
#include "Recorder.hpp"

/* Create a Recorder instance. Subscribe to command input and set audio callback.
 * Only will be active if enabled by the configuration, in which case frames will
 * automatically begin to be stashed */
Recorder::Recorder(Node* parent) : Node(parent)
{
    get_delegate().subscribe(&Recorder::respond, this);
    animation.play();
    Mix_SetPostMix(Recorder::process_audio, this);
    if (!configuration()["recording"]["enabled"])
    {
        deactivate();
    }
}

/* Returns length of a recorded video frame in seconds. Defaults to the frame length of the game if this hasn't
 * been configured by the user. */
float Recorder::frame_length()
{
    return configuration()["recording"].value("video-frame-length", get_root()->get_frame_length());
}

/* Handle commands for screenshot, record video and save video */
void Recorder::respond(SDL_Event& event)
{
    if (get_delegate().compare(event, "screenshot"))
    {
        capture_screen();
    }
    else if (get_delegate().compare(event, "record"))
    {
        if (is_recording)
        {
            end_recording();
        }
        else if (!writing_recording)
        {
            start_recording();
        }
        else
        {
            sb::Log::log("Writing in progress, cannot start recording");
        }
    }
    else if (get_delegate().compare(event, "save-current-stash"))
    {
        grab_stash();
    }
    else if (get_delegate().compare(event, "print-video-memory-size"))
    {
        std::ostringstream message;
        message << "Video memory size is " << get_memory_size() << "MB";
        sb::Log::log(message);
    }
}

/* Save the current screen pixels as a PNG in the path specified by the configuration. Parent
 * directories in the path will be created if necessary. The file name will be auto generated
 * based on the configuration. An index will be included in the file name based on previous
 * screenshots found in the output directory. */
void Recorder::capture_screen()
{
    nlohmann::json config = configuration();
    SDL_Surface* surface = get_display().screen_surface();
    fs::path directory = config["recording"]["screenshot-directory"];
    fs::create_directories(directory);
    std::string prefix = config["recording"]["screenshot-prefix"].get<std::string>();
    std::string extension = config["recording"]["screenshot-extension"].get<std::string>();
    int zfill = config["recording"]["screenshot-zfill"];
    fs::path path = sb::get_next_file_name(directory, zfill, prefix, extension);
    IMG_SavePNG(surface, path.c_str());
    SDL_FreeSurface(surface);
    std::ostringstream message;
    message << "saved screenshot to " << path;
    sb::Log::log(message);
}

/* Writes a video of what was just displayed on the screen up until the function was called. The length
 * of the video is determined by the stash length written in the configuration. This is accomplished by
 * writing the contents of the most recent Stash object. */
void Recorder::grab_stash()
{
    if (!is_recording and !writing_recording)
    {
        int length = configuration()["recording"]["max-stash-length"];
        std::ostringstream message;
        message << "stashing most recent " << length / 1000.0f << " seconds of video";
        sb::Log::log(message);
        most_recent_stash = current_stash;
        current_stash = Stash();
        writing_recording = true;
        std::function<void()> f = std::bind(&Recorder::write_most_recent_frames, this);
        std::thread writing(f);
        writing.detach();
    }
    else
    {
        sb::Log::log("recording in progress, cannot grab most recent frames");
    }
}

void Recorder::write_most_recent_frames()
{
    make_directory();
    write_stash_frames(&most_recent_stash);
    open_audio_file();
    while (!most_recent_stash.audio_buffers.empty())
    {
        write_audio(most_recent_stash.audio_buffers.front(), most_recent_stash.audio_buffer_lengths.front());
        most_recent_stash.audio_buffers.erase(most_recent_stash.audio_buffers.begin());
        most_recent_stash.audio_buffer_lengths.erase(most_recent_stash.audio_buffer_lengths.begin());
    }
    audio_file.close();
    if (configuration()["recording"]["write-mp4"])
    {
        write_mp4();
    }
    std::ostringstream message;
    message << "wrote video frames to " << current_video_directory;
    sb::Log::log(message);
    writing_recording = false;
}

void Recorder::start_recording()
{
    if (!writing_recording)
    {
        sb::Log::log("starting recording");
        is_recording = true;
        video_stashes.push_back(Stash());
        make_directory();
        open_audio_file();
    }
    else
    {
        sb::Log::log("writing in progress, cannot start recording", sb::Log::WARN);
    }
}

void Recorder::open_audio_file()
{
    std::stringstream audio_path;
    audio_path << current_video_directory.string() << ".raw";
    current_audio_path = audio_path.str();
    audio_file.open(audio_path.str(), std::ios::binary);
}

void Recorder::add_frame()
{
    glm::ivec2 size = get_display().window_size();
    int bytes = sb::Display::bpp / 8 * size.x * size.y;
    unsigned char* pixels = new unsigned char[bytes];
    get_display().screen_pixels(pixels, size.x, size.y);
    int max_length = configuration()["recording"]["max-stash-length"];
    float length = frame_length() * current_stash.pixel_buffers.size();
    if (length > max_length)
    {
        delete[] current_stash.pixel_buffers.front();
        current_stash.pixel_buffers.erase(current_stash.pixel_buffers.begin());
        current_stash.flipped.erase(current_stash.flipped.begin());
    }
    current_stash.pixel_buffers.push_back(pixels);
    current_stash.flipped.push_back(get_root()->is_gl_context);
    if (is_recording)
    {
        unsigned char* vid_pixels = new unsigned char[bytes];
        memcpy(vid_pixels, pixels, bytes);
        video_stashes.back().pixel_buffers.push_back(vid_pixels);
        video_stashes.back().flipped.push_back(get_root()->is_gl_context);
        if (video_stashes.back().pixel_buffers.size() * frame_length() > max_length)
        {
            std::function<void(Stash*)> f = std::bind(&Recorder::write_stash_frames, this, std::placeholders::_1);
            std::thread writing(f, &video_stashes.back());
            writing.detach();
            // int frame_offset = video_stashes.back().frame_offset;
            // std::vector<unsigned char*> pixel_buffers = video_stashes.back().pixel_buffers;
            // std::vector<bool> flipped = video_stashes.back().flipped;
            // for (int ii = frame_offset; ii < pixel_buffers.size() + frame_offset; ii++)
            // {
            //     SDL_Surface* frame = get_display().screen_surface_from_pixels(
            //         pixel_buffers[ii - frame_offset], flipped[ii - frame_offset]);
            //     std::stringstream name;
            //     name << sb::pad(ii, 5) << ".png";
            //     fs::path path = current_video_directory / name.str();
            //     SDL_Log("%s (%i, %i) (%i, %i, %i, %i)", path.c_str(), frame->w, frame->h,
            //             ((unsigned char*) frame->pixels)[0], ((unsigned char*) frame->pixels)[1],
            //             ((unsigned char*) frame->pixels)[2], ((unsigned char*) frame->pixels)[3]);
            //     IMG_SavePNG(frame, path.string().c_str());
            // }
            // end_recording();
            // write_stash_frames(video_stashes.back().pixel_buffers,
            //                    video_stashes.back().flipped,
            //                    video_stashes.back().frame_offset);
            video_stashes.push_back(Stash(video_stashes.back().frame_offset +
                                          video_stashes.back().pixel_buffers.size()));
        }
    }
}

int Recorder::get_memory_size()
{
    glm::ivec2 window = get_display().window_size();
    int bytes_per_frame = sb::Display::bpp / 8 * window.x * window.y,
        size_in_bytes = 0;
    for (Stash& stash : in_game_stashes)
    {
        size_in_bytes += stash.pixel_buffers.size() * bytes_per_frame;
        for (int& length : stash.audio_buffer_lengths)
        {
            size_in_bytes += length;
        }
    }
    for (Stash& stash : video_stashes)
    {
        size_in_bytes += stash.pixel_buffers.size() * bytes_per_frame;
    }
    size_in_bytes += current_stash.pixel_buffers.size() * bytes_per_frame;
    for (int& length : current_stash.audio_buffer_lengths)
    {
        size_in_bytes += length;
    }
    size_in_bytes += most_recent_stash.pixel_buffers.size() * bytes_per_frame;
    for (int& length : most_recent_stash.audio_buffer_lengths)
    {
        size_in_bytes += length;
    }
    return size_in_bytes / 1000000;
}

void Recorder::make_directory()
{
    nlohmann::json config = configuration();
    fs::path root = config["recording"]["video-directory"];
    fs::create_directories(root);
    fs::path directory = sb::get_next_file_name(root, 5, "video-");
    fs::create_directories(directory);
    current_video_directory = directory;
}

void Recorder::write_stash_frames(Stash* stash)
{
    SDL_Log("Writing stash offset %i to %s...", stash->frame_offset, current_video_directory.c_str());
    SDL_Surface* frame;
    GifWriter gif_writer;
    int gif_frame_length = configuration()["recording"]["gif-frame-length"];
    fs::path gif_path = sb::get_next_file_name(
        current_video_directory, 3, "gif-", ".gif");
    float elapsed = 0, last_gif_write = 0, gif_write_overflow = 0;
    for (int ii = stash->frame_offset; not stash->pixel_buffers.empty(); ii++)
    {
        frame = get_display().screen_surface_from_pixels(
            stash->pixel_buffers.front(), stash->flipped.front());
        std::stringstream name;
        name << sb::pad(ii, 5) << ".png";
        fs::path path = current_video_directory / name.str();
        IMG_SavePNG(frame, path.string().c_str());
        if (ii == stash->frame_offset or
            elapsed - last_gif_write + gif_write_overflow >= gif_frame_length)
        {
            if (ii == stash->frame_offset)
            {
                GifBegin(&gif_writer, gif_path.string().c_str(), frame->w,
                         frame->h, gif_frame_length / 10);
            }
            else
            {
                gif_write_overflow += elapsed - (last_gif_write + gif_frame_length);
                last_gif_write = elapsed;
            }
            SDL_Surface* converted = SDL_ConvertSurfaceFormat(
                frame, SDL_PIXELFORMAT_ABGR8888, 0);
            GifWriteFrame(&gif_writer, (const uint8_t*) converted->pixels,
                          frame->w, frame->h, gif_frame_length / 10);
        }
        elapsed += frame_length();
        delete[] stash->pixel_buffers.front();
        stash->pixel_buffers.erase(stash->pixel_buffers.begin());
        stash->flipped.erase(stash->flipped.begin());
        SDL_FreeSurface(frame);
    }
    GifEnd(&gif_writer);
}

void Recorder::keep_stash()
{
    in_game_stashes.push_back(current_stash);
    current_stash = Stash();
    auto max_stashes = configuration()["recording"]["max-in-game-stashes"];
    if (in_game_stashes.size() > max_stashes)
    {
        Stash& stash = in_game_stashes.front();
        while (not stash.pixel_buffers.empty())
        {
            delete[] stash.pixel_buffers.back();
            stash.pixel_buffers.pop_back();
            stash.flipped.pop_back();
        }
        in_game_stashes.erase(in_game_stashes.begin());
    }
}

void Recorder::end_recording()
{
    std::cout << "Ending recording..." << std::endl;
    audio_file.close();
    is_recording = false;
    writing_recording = true;
    std::function<void()> f = std::bind(&Recorder::finish_writing_video, this);
    std::thread finishing(f);
    finishing.detach();
}

void Recorder::finish_writing_video()
{
    write_stash_frames(&video_stashes.back());
    int count;
    while (true)
    {
        count = 0;
        for (Stash& stash : video_stashes)
        {
            count += stash.pixel_buffers.size();
        }
        if (count == 0)
        {
            break;
        }
    }
    video_stashes.clear();
    if (configuration()["recording"]["write-mp4"])
    {
        write_mp4();
    }
    std::cout << "Wrote video frames to " << current_video_directory.string() << std::endl;
    writing_recording = false;
}

/* Launch a system command that calls ffmpeg to write an x264 encoded MP4 from the image frames written.
 * This requires ffmpeg to be installed on the user's system. Might only work on Linux (?) */
void Recorder::write_mp4()
{
    glm::ivec2 size = get_display().window_size();
    std::ostringstream mp4_command;
    std::string pixel_format = configuration()["recording"]["mp4-pixel-format"].get<std::string>();
    fs::path images_match = current_video_directory / "%05d.png";
    mp4_command << "ffmpeg -f s16le -ac 2 -ar 22050 -i " << current_audio_path.string() <<
        " -f image2 -framerate " << (1000 / frame_length()) <<
        " -i " << images_match.string() << " -s " << size.x << "x" << size.y <<
        " -c:v libx264 -crf 17 -pix_fmt " << pixel_format << " " <<
        current_video_directory.string() << ".mp4";
    std::string mp4_command_str = mp4_command.str();
    std::cout << mp4_command_str << std::endl;
    std::system(mp4_command_str.c_str());
}

void Recorder::write_audio(Uint8* stream, int len)
{
    audio_file.write(reinterpret_cast<char*>(stream), len);
}

void Recorder::update()
{
    if (is_recording and get_memory_size() > configuration()["recording"]["max-video-memory"])
    {
        end_recording();
    }
    animation.set_frame_length(frame_length());
    animation.update();
}

void Recorder::process_audio(void* context, Uint8* stream, int len)
{
    Recorder* recorder = static_cast<Recorder*>(context);
    if (recorder->is_active())
    {
        int max_length = recorder->configuration()["recording"]["max-stash-length"];
        float length = recorder->frame_length() * recorder->current_stash.pixel_buffers.size();
        if (length > max_length)
        {
            delete[] recorder->current_stash.audio_buffers.front();
            recorder->current_stash.audio_buffers.erase(recorder->current_stash.audio_buffers.begin());
            recorder->current_stash.audio_buffer_lengths.erase(
                recorder->current_stash.audio_buffer_lengths.begin());
        }
        Uint8* stream_copy = new Uint8[len];
        std::memcpy(stream_copy, stream, len);
        recorder->current_stash.audio_buffers.push_back(stream_copy);
        recorder->current_stash.audio_buffer_lengths.push_back(len);
        if (recorder->is_recording)
        {
            recorder->write_audio(stream_copy, len);
        }
    }
}