#include "media/hardware_decode.hpp"

#ifdef SMARTCUT_HAS_FFMPEG
extern "C" {
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
#include <libavutil/hwcontext.h>
}
#endif

#include <array>
#include <cstdlib>
#include <string_view>

namespace smartcut::media {

#ifdef SMARTCUT_HAS_FFMPEG

HardwareVideoDecoder::~HardwareVideoDecoder() { close(); }

void HardwareVideoDecoder::close() {
  avcodec_free_context(&context_);
  av_buffer_unref(&hardware_device_);
  hardware_format_ = AV_PIX_FMT_NONE;
  hardware_requested_ = false;
  hardware_frame_seen_ = false;
  backend_name_ = "software";
}

AVPixelFormat HardwareVideoDecoder::select_pixel_format(
    AVCodecContext* context, const AVPixelFormat* formats) {
  const auto* session = static_cast<const HardwareVideoDecoder*>(context->opaque);
  if (session && session->hardware_format_ != AV_PIX_FMT_NONE) {
    for (const auto* format = formats; *format != AV_PIX_FMT_NONE; ++format)
      if (*format == session->hardware_format_) return *format;
  }
  return avcodec_default_get_format(context, formats);
}

bool HardwareVideoDecoder::open_software(const AVCodecParameters* parameters,
                                         const AVCodec* decoder) {
  if (!decoder) return false;
  context_ = avcodec_alloc_context3(decoder);
  if (!context_ || avcodec_parameters_to_context(context_, parameters) < 0 ||
      avcodec_open2(context_, decoder, nullptr) < 0) {
    avcodec_free_context(&context_);
    return false;
  }
  backend_name_ = "software";
  return true;
}

bool HardwareVideoDecoder::open(const AVCodecParameters* parameters,
                                AVCodecID codec_id,
                                const AVCodec* software_decoder,
                                bool allow_hardware) {
  close();
  if (const auto* disabled = std::getenv("SMARTCUT_DISABLE_HARDWARE_DECODE")) {
    const std::string_view value(disabled);
    if (value == "1" || value == "true" || value == "TRUE") allow_hardware = false;
  }
#ifdef _WIN32
  if (allow_hardware) {
    const auto* hardware_decoder = avcodec_find_decoder(codec_id);
    if (hardware_decoder) {
      constexpr std::array preferred_devices{AV_HWDEVICE_TYPE_D3D11VA,
                                              AV_HWDEVICE_TYPE_VULKAN};
      for (const auto device_type : preferred_devices) {
        const AVCodecHWConfig* configuration = nullptr;
        hardware_format_ = AV_PIX_FMT_NONE;
        for (int index = 0;
             (configuration = avcodec_get_hw_config(hardware_decoder, index)); ++index) {
          if (configuration->device_type == device_type &&
              (configuration->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) != 0) {
            hardware_format_ = configuration->pix_fmt;
            break;
          }
        }
        if (hardware_format_ == AV_PIX_FMT_NONE ||
            av_hwdevice_ctx_create(&hardware_device_, device_type, nullptr, nullptr, 0) < 0)
          continue;
        context_ = avcodec_alloc_context3(hardware_decoder);
        if (context_ && avcodec_parameters_to_context(context_, parameters) >= 0) {
          context_->opaque = this;
          context_->get_format = &HardwareVideoDecoder::select_pixel_format;
          context_->hw_device_ctx = av_buffer_ref(hardware_device_);
          if (context_->hw_device_ctx &&
              avcodec_open2(context_, hardware_decoder, nullptr) >= 0) {
            hardware_requested_ = true;
            const char* name = av_hwdevice_get_type_name(device_type);
            backend_name_ = name ? name : "hardware";
            return true;
          }
        }
        avcodec_free_context(&context_);
        av_buffer_unref(&hardware_device_);
      }
    }
    av_buffer_unref(&hardware_device_);
    hardware_format_ = AV_PIX_FMT_NONE;
  }
#else
  (void)allow_hardware;
  (void)codec_id;
#endif
  return open_software(parameters, software_decoder);
}

const AVFrame* HardwareVideoDecoder::cpu_frame(const AVFrame* decoded_frame,
                                                AVFrame* transfer_frame) {
  if (!decoded_frame) return nullptr;
  if (decoded_frame->format != hardware_format_ || hardware_format_ == AV_PIX_FMT_NONE)
    return decoded_frame;
  hardware_frame_seen_ = true;
  if (!transfer_frame) return nullptr;
  av_frame_unref(transfer_frame);
  if (av_hwframe_transfer_data(transfer_frame, decoded_frame, 0) < 0) return nullptr;
  if (av_frame_copy_props(transfer_frame, decoded_frame) < 0) return nullptr;
  return transfer_frame;
}

bool HardwareVideoDecoder::is_hardware_frame(const AVFrame* frame) const {
  return frame && hardware_format_ != AV_PIX_FMT_NONE && frame->format == hardware_format_;
}

HardwareFrameScaler::~HardwareFrameScaler() { reset(); }

void HardwareFrameScaler::reset() {
  avfilter_graph_free(&graph_);
  source_ = nullptr;
  scaler_ = nullptr;
  sink_ = nullptr;
  source_width_ = source_height_ = output_width_ = output_height_ = 0;
  source_format_ = AV_PIX_FMT_NONE;
  backend_name_.clear();
}

bool HardwareFrameScaler::configure(const AVFrame* source, int width, int height) {
  reset();
  if (!source || !source->hw_frames_ctx || width <= 0 || height <= 0) return false;
  const auto source_format = static_cast<AVPixelFormat>(source->format);
  const char* scale_name = nullptr;
  std::string scale_arguments;
  if (source_format == AV_PIX_FMT_D3D11) {
    scale_name = "scale_d3d11";
    scale_arguments = "width=" + std::to_string(width) + ":height=" +
                      std::to_string(height) + ":format=nv12";
    backend_name_ = "scale_d3d11";
  } else if (source_format == AV_PIX_FMT_VULKAN) {
    scale_name = "scale_vulkan";
    scale_arguments = "w=" + std::to_string(width) + ":h=" + std::to_string(height);
    backend_name_ = "scale_vulkan";
  } else {
    return false;
  }
  graph_ = avfilter_graph_alloc();
  const auto* buffer = avfilter_get_by_name("buffer");
  const auto* scale = avfilter_get_by_name(scale_name);
  const auto* sink = avfilter_get_by_name("buffersink");
  if (!graph_ || !buffer || !scale || !sink) {
    reset();
    return false;
  }
  source_ = avfilter_graph_alloc_filter(graph_, buffer, "hardware-input");
  int result = source_ ? 0 : AVERROR(ENOMEM);
  if (source_) {
    auto* parameters = av_buffersrc_parameters_alloc();
    if (!parameters) result = AVERROR(ENOMEM);
    else {
      parameters->format = source_format;
      parameters->width = source->width;
      parameters->height = source->height;
      parameters->time_base = AVRational{1, 1000000};
      parameters->sample_aspect_ratio = source->sample_aspect_ratio.num > 0
                                            ? source->sample_aspect_ratio : AVRational{1, 1};
      parameters->color_space = source->colorspace;
      parameters->color_range = source->color_range;
      parameters->alpha_mode = source->alpha_mode;
      parameters->hw_frames_ctx = av_buffer_ref(source->hw_frames_ctx);
      result = parameters->hw_frames_ctx
                   ? av_buffersrc_parameters_set(source_, parameters) : AVERROR(ENOMEM);
      av_buffer_unref(&parameters->hw_frames_ctx);
      av_freep(&parameters);
    }
    if (result >= 0) result = avfilter_init_dict(source_, nullptr);
  }
  if (result >= 0)
    result = avfilter_graph_create_filter(&scaler_, scale, "hardware-scale",
                                           scale_arguments.c_str(), nullptr, graph_);
  if (result >= 0)
    result = avfilter_graph_create_filter(&sink_, sink, "hardware-output", nullptr, nullptr, graph_);
  if (result >= 0) result = avfilter_link(source_, 0, scaler_, 0);
  if (result >= 0) result = avfilter_link(scaler_, 0, sink_, 0);
  if (result >= 0) result = avfilter_graph_config(graph_, nullptr);
  if (result < 0) {
    reset();
    return false;
  }
  source_width_ = source->width;
  source_height_ = source->height;
  output_width_ = width;
  output_height_ = height;
  source_format_ = source_format;
  return true;
}

const AVFrame* HardwareFrameScaler::scale(const AVFrame* source, AVFrame* scaled,
                                          int width, int height) {
  if (!source || !scaled ||
      (source->format != AV_PIX_FMT_D3D11 && source->format != AV_PIX_FMT_VULKAN))
    return nullptr;
  if (!graph_ || source_width_ != source->width || source_height_ != source->height ||
      output_width_ != width || output_height_ != height ||
      source_format_ != static_cast<AVPixelFormat>(source->format)) {
    if (!configure(source, width, height)) return nullptr;
  }
  av_frame_unref(scaled);
  if (av_buffersrc_add_frame_flags(source_, const_cast<AVFrame*>(source),
                                   AV_BUFFERSRC_FLAG_KEEP_REF) < 0 ||
      av_buffersink_get_frame(sink_, scaled) < 0) return nullptr;
  return scaled;
}

#endif

}  // namespace smartcut::media
