/***

  Olive - Non-Linear Video Editor
  Copyright (C) 2022 Olive Team

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.

***/

#include "export.h"

#include "node/color/colormanager/colormanager.h"

namespace olive {

ExportTask::ExportTask(ViewerOutput *viewer_node,
                       ColorManager* color_manager,
                       const EncodingParams& params) :
  params_(params)
{
  // Create a copy of the project
  copier_ = new ProjectCopier(this);
  copier_->SetProject(viewer_node->project());

  set_viewer(copier_->GetCopy(viewer_node));
  color_manager_ = copier_->GetCopiedProject()->color_manager();

  // Adjust video params to have no divider
  VideoParams vp = viewer_node->GetVideoParams();
  vp.set_divider(1);
  vp.set_time_base(params.video_params().time_base());
  vp.set_frame_rate(params.video_params().frame_rate());
  set_video_params(vp);

  set_audio_params(viewer_node->GetAudioParams());

  SetTitle(tr("Exporting \"%1\"").arg(viewer_node->GetLabel()));
  SetNativeProgressSignallingEnabled(false);
}

bool ExportTask::Run()
{
  // For safety, if we're overwriting, we save to a temporary filename and then only overwrite it
  // at the end
  QString real_filename = params_.filename();
  if (QFileInfo::exists(params_.filename())) {
    // Generate a filename that definitely doesn't exist
    params_.SetFilename(FileFunctions::GetSafeTemporaryFilename(real_filename));
  }

  // If we're exporting to a sidecar subtitle file, disable the subtitles in the main encoder
  bool subtitles_enabled = params_.subtitles_enabled();
  EncodingParams sidecar_params = params_;
  if (subtitles_enabled && params_.subtitles_are_sidecar()) {
    params_.DisableSubtitles();
  }

  encoder_ = std::shared_ptr<Encoder>(Encoder::CreateFromParams(params_));

  if (!encoder_) {
    SetError(tr("Failed to create encoder"));
    return false;
  }

  if (!encoder_->Open()) {
    SetError(tr("Failed to open file: %1").arg(encoder_->GetError()));
    return false;
  }

  if (subtitles_enabled && params_.subtitles_are_sidecar()) {
    // Construct sidecar params
    sidecar_params.DisableVideo();
    sidecar_params.DisableAudio();

    QString sidecar_filename;
    {
      QFileInfo fi(real_filename);
      sidecar_filename = fi.completeBaseName();
      sidecar_filename.append('.');
      sidecar_filename.append(ExportFormat::GetExtension(sidecar_params.subtitle_sidecar_fmt()));
      sidecar_filename = fi.dir().filePath(sidecar_filename);
    }
    sidecar_params.SetFilename(sidecar_filename);

    subtitle_encoder_ = std::shared_ptr<Encoder>(Encoder::CreateFromFormat(sidecar_params.subtitle_sidecar_fmt(), sidecar_params));
    if (!subtitle_encoder_) {
      SetError(tr("Failed to create subtitle encoder"));
      return false;
    }

    if (!subtitle_encoder_->Open()) {
      SetError(tr("Failed to open subtitle sidecar file: %1").arg(sidecar_filename));
      return false;
    }
  } else {
    subtitle_encoder_ = encoder_;
  }

  if (params_.has_custom_range()) {
    // Render custom range only
    export_range_ = params_.custom_range();
  } else {
    // Render entire sequence
    export_range_ = TimeRange(0, viewer()->GetLength());
  }

  frame_time_ = 0;

  QSize video_force_size;
  QMatrix4x4 video_force_matrix;

  if (params_.video_enabled()) {

    // If a transformation matrix is applied to this video, create it here
    if (video_params().width() != params_.video_params().width()
        || video_params().height() != params_.video_params().height()) {
      video_force_size = QSize(params_.video_params().width(), params_.video_params().height());

      if (params_.video_scaling_method() != EncodingParams::kStretch) {
        video_force_matrix = EncodingParams::GenerateMatrix(params_.video_scaling_method(),
                                                            video_params().width(),
                                                            video_params().height(),
                                                            params_.video_params().width(),
                                                            params_.video_params().height());
      }
    } else {
      // Disables forcing size in the renderer
      video_force_size = QSize(0, 0);
    }

    // Create color processor
    color_processor_ = ColorProcessor::Create(color_manager_,
                                              color_manager_->GetReferenceColorSpace(),
                                              params_.color_transform());
  }

  // Start render process
  TimeRangeList video_range, audio_range;
  TimeRange subtitle_range;

  if (params_.video_enabled()) {
    if (export_range_.in() > 0) {
      export_range_.set_in(Timecode::snap_time_to_timebase(export_range_.in(), video_params().frame_rate_as_time_base()));
    }

    video_range = {export_range_};
  }

  if (params_.audio_enabled()) {
    audio_range = {export_range_};
  }

  if (subtitles_enabled) {
    subtitle_range = export_range_;
  }

  Render(color_manager_, video_range, audio_range, subtitle_range, RenderMode::kOnline, nullptr,
         video_force_size, video_force_matrix, encoder_->GetDesiredPixelFormat(),
         VideoParams::kRGBAChannelCount, color_processor_);

  bool success = true;

  encoder_->Close();
  if (!encoder_->GetError().isEmpty()) {
    SetError(encoder_->GetError());
    success = false;
  }

  if (subtitle_encoder_ != encoder_) {
    subtitle_encoder_->Close();
    if (!subtitle_encoder_->GetError().isEmpty()) {
      SetError(subtitle_encoder_->GetError());
      success = false;
    }
  }

  // If cancelled, delete the file we made, which is always a file we created since we write to a
  // temp file during the actual encoding process
  if (IsCancelled()) {
    QFile::remove(params_.filename());
  } else if (params_.filename() != real_filename) {
    // If we were writing to a temp file, overwrite now
    if (!FileFunctions::RenameFileAllowOverwrite(params_.filename(), real_filename)) {
      SetError(tr("Failed to overwrite \"%1\". Export has been saved as \"%2\" instead.")
               .arg(real_filename, params_.filename()));
      success = false;
    }
  }

  return success;
}

bool ExportTask::FrameDownloaded(FramePtr f, const rational &time)
{
  rational actual_time = time - export_range_.in();

  time_map_.insert(actual_time, f);

  while (!IsCancelled()) {
    rational real_time = Timecode::timestamp_to_time(frame_time_,
                                                     video_params().frame_rate_as_time_base());

    if (!time_map_.contains(real_time)) {
      break;
    }

    // Unfortunately this can't be done in another thread since the frames need to be sent
    // one after the other chronologically.
    if (!encoder_->WriteFrame(time_map_.take(real_time), real_time)) {
      SetError(encoder_->GetError());
      return false;
    }

    frame_time_++;
    emit ProgressChanged(double(frame_time_) / double(GetTotalNumberOfFrames()));
  }

  return true;
}

bool ExportTask::AudioDownloaded(const TimeRange &range, const SampleBuffer &samples)
{
  TimeRange adjusted_range = range - export_range_.in();

  if (adjusted_range.in() == audio_time_) {
    if (!WriteAudioLoop(adjusted_range, samples)) {
      return false;
    }
  } else {
    audio_map_.insert(adjusted_range, samples);
  }

  return true;
}

bool ExportTask::EncodeSubtitle(const SubtitleBlock *sub)
{
  if (!subtitle_encoder_->WriteSubtitle(sub)) {
    SetError(subtitle_encoder_->GetError());
    return false;
  } else {
    return true;
  }
}

bool ExportTask::WriteAudioLoop(const TimeRange& time, const SampleBuffer &samples)
{
  if (!encoder_->WriteAudio(samples)) {
    SetError(encoder_->GetError());
    return false;
  }

  audio_time_ = time.out();

  for (auto it=audio_map_.begin(); it!=audio_map_.end(); it++) {
    TimeRange t = it.key();
    SampleBuffer s = it.value();

    if (t.in() == audio_time_) {
      // Erase from audio map since we're just about to write it
      audio_map_.erase(it);

      // Call recursively to write the next sample buffer
      if (!WriteAudioLoop(t, s)) {
        return false;
      }

      // Break out of loop
      break;
    }
  }

  return true;
}

}