#include "media/title_media.hpp"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>

#include "media/ffmpeg_raii.hpp"

#ifdef SMARTCUT_HAS_FFMPEG
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
}
#endif

namespace smartcut::media {
namespace {
#ifdef SMARTCUT_HAS_FFMPEG
int primary_video_stream(AVFormatContext* context) {
  for (unsigned index = 0; index < context->nb_streams; ++index) {
    const auto* parameters = context->streams[index]->codecpar;
    if (parameters->codec_type == AVMEDIA_TYPE_VIDEO &&
        (context->streams[index]->disposition & AV_DISPOSITION_ATTACHED_PIC) == 0) return static_cast<int>(index);
  }
  return -1;
}

bool encoder_supports(const AVCodec* codec, AVPixelFormat format) {
  const void* configurations = nullptr; int count = 0;
  if (avcodec_get_supported_config(nullptr, codec, AV_CODEC_CONFIG_PIX_FORMAT, 0, &configurations, &count) < 0 || !configurations) return true;
  const auto* formats = static_cast<const AVPixelFormat*>(configurations);
  return std::find(formats, formats + count, format) != formats + count;
}

AVPixelFormat select_pixel_format(const AVCodec* codec, AVPixelFormat requested) {
  if (requested != AV_PIX_FMT_NONE && encoder_supports(codec, requested)) return requested;
  if (requested == AV_PIX_FMT_YUV420P10LE && encoder_supports(codec, AV_PIX_FMT_YUV420P10LE)) return AV_PIX_FMT_YUV420P10LE;
  if (encoder_supports(codec, AV_PIX_FMT_YUV420P)) return AV_PIX_FMT_YUV420P;
  const void* configurations = nullptr; int count = 0;
  if (avcodec_get_supported_config(nullptr, codec, AV_CODEC_CONFIG_PIX_FORMAT, 0, &configurations, &count) >= 0 && configurations && count > 0)
    return static_cast<const AVPixelFormat*>(configurations)[0];
  return AV_PIX_FMT_YUV420P;
}
#endif
}  // namespace

std::string create_title_media(const std::filesystem::path& reference,
                               const std::filesystem::path& bgra_image,
                               int image_width, int image_height,
                               const std::filesystem::path& output,
                               double duration_seconds,
                               double fade_in_seconds,
                               double fade_out_seconds,
                               const std::string& description) {
#ifndef SMARTCUT_HAS_FFMPEG
  (void)reference; (void)bgra_image; (void)image_width; (void)image_height; (void)output;
  (void)duration_seconds; (void)fade_in_seconds; (void)fade_out_seconds; (void)description;
  return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"FFmpeg unavailable\"}";
#else
  if (image_width <= 0 || image_height <= 0 || duration_seconds <= 0.0 || fade_in_seconds < 0.0 || fade_out_seconds < 0.0)
    return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"invalid title parameters\"}";
  std::error_code path_error;
  if (std::filesystem::weakly_canonical(reference, path_error) == std::filesystem::weakly_canonical(output, path_error))
    return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"reference and output must differ\"}";
  const auto byte_count = static_cast<std::uint64_t>(image_width) * static_cast<std::uint64_t>(image_height) * 4ULL;
  if (byte_count > 1024ULL * 1024ULL * 1024ULL)
    return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title image is too large\"}";
  std::ifstream image_file(bgra_image, std::ios::binary);
  std::vector<std::uint8_t> original(static_cast<std::size_t>(byte_count));
  if (!image_file || !image_file.read(reinterpret_cast<char*>(original.data()), static_cast<std::streamsize>(original.size())) || image_file.peek() != std::ifstream::traits_type::eof())
    return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"BGRA title image size mismatch\"}";

  AVFormatContext* source = nullptr;
  AVFormatContext* destination = nullptr;
  AVCodecContext* encoder = nullptr;
  SwsContext* scaler = nullptr;
  AVFrame* frame = av_frame_alloc();
  AVPacket* packet = av_packet_alloc();
  bool completed = false;
  auto cleanup = [&] {
    av_packet_free(&packet); av_frame_free(&frame); sws_freeContext(scaler); avcodec_free_context(&encoder);
    if (destination && destination->pb) avio_closep(&destination->pb);
    avformat_free_context(destination); avformat_close_input(&source);
    if (!completed) { std::error_code ignored; std::filesystem::remove(output, ignored); }
  };
  if (!frame || !packet || avformat_open_input(&source, reference.string().c_str(), nullptr, nullptr) < 0 ||
      avformat_find_stream_info(source, nullptr) < 0) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"reference open failed\"}"; }
  const int video_index = primary_video_stream(source);
  if (video_index < 0) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"reference video stream missing\"}"; }
  const auto* source_stream = source->streams[video_index];
  const auto* source_parameters = source_stream->codecpar;
  const auto source_stream_count = source->nb_streams;
  const auto* codec = ffmpeg::matching_boundary_encoder(source_parameters->codec_id);
  if (!codec) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"matching title encoder unavailable\"}"; }
  if (avformat_alloc_output_context2(&destination, nullptr, "matroska", output.string().c_str()) < 0 || !destination) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"output setup failed\"}"; }
  encoder = avcodec_alloc_context3(codec);
  if (!encoder) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"encoder allocation failed\"}"; }
  AVRational rate = source_stream->avg_frame_rate.num > 0 && source_stream->avg_frame_rate.den > 0 ? source_stream->avg_frame_rate : source_stream->r_frame_rate;
  if (rate.num <= 0 || rate.den <= 0) rate = AVRational{25, 1};
  const auto frames = std::max<std::int64_t>(1, static_cast<std::int64_t>(std::llround(duration_seconds * av_q2d(rate))));
  encoder->width = source_parameters->width; encoder->height = source_parameters->height;
  encoder->pix_fmt = select_pixel_format(codec, static_cast<AVPixelFormat>(source_parameters->format));
  encoder->framerate = rate; encoder->time_base = av_inv_q(rate);
  encoder->sample_aspect_ratio = source_parameters->sample_aspect_ratio.num > 0 ? source_parameters->sample_aspect_ratio : AVRational{1, 1};
  encoder->color_range = source_parameters->color_range; encoder->color_primaries = source_parameters->color_primaries;
  encoder->color_trc = source_parameters->color_trc; encoder->colorspace = source_parameters->color_space;
  encoder->chroma_sample_location = source_parameters->chroma_location;
  encoder->bit_rate = source_parameters->bit_rate > 0 ? source_parameters->bit_rate : 8'000'000;
  encoder->profile = source_parameters->profile; encoder->level = source_parameters->level;
  encoder->gop_size = static_cast<int>(std::min<std::int64_t>(frames, 100000)); encoder->max_b_frames = 0; encoder->thread_count = 0;
  if (destination->oformat->flags & AVFMT_GLOBALHEADER) encoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  const auto repeat = std::string("repeat-headers=1:keyint=") + std::to_string(frames) + ":min-keyint=" + std::to_string(frames) + ":scenecut=0";
  if (source_parameters->codec_id == AV_CODEC_ID_H264) av_opt_set(encoder->priv_data, "x264-params", repeat.c_str(), 0);
  else if (source_parameters->codec_id == AV_CODEC_ID_HEVC) {
    av_opt_set(encoder->priv_data, "x265-params", repeat.c_str(), 0);
    av_opt_set(encoder->priv_data, "profile", encoder->pix_fmt == AV_PIX_FMT_YUV420P10LE ? "main10" : "main", 0);
  } else if (source_parameters->codec_id == AV_CODEC_ID_AV1 || source_parameters->codec_id == AV_CODEC_ID_VP9) {
    av_opt_set(encoder->priv_data, "cpu-used", "8", 0); av_opt_set(encoder->priv_data, "row-mt", "1", 0);
    av_opt_set(encoder->priv_data, "lag-in-frames", "0", 0);
    if (source_parameters->codec_id == AV_CODEC_ID_VP9) { av_opt_set(encoder->priv_data, "deadline", "realtime", 0); av_opt_set(encoder->priv_data, "auto-alt-ref", "0", 0); }
  }
  if (avcodec_open2(encoder, codec, nullptr) < 0) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title encoder open failed\"}"; }
  for (unsigned stream_index = 0; stream_index < source->nb_streams; ++stream_index) {
    auto* output_stream = avformat_new_stream(destination, nullptr);
    if (!output_stream) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title stream allocation failed\"}"; }
    const int copy_result = static_cast<int>(stream_index) == video_index ? avcodec_parameters_from_context(output_stream->codecpar, encoder) : avcodec_parameters_copy(output_stream->codecpar, source->streams[stream_index]->codecpar);
    if (copy_result < 0) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title stream setup failed\"}"; }
    output_stream->codecpar->codec_tag = 0;
    output_stream->time_base = static_cast<int>(stream_index) == video_index ? encoder->time_base : source->streams[stream_index]->time_base;
    output_stream->disposition = source->streams[stream_index]->disposition;
    output_stream->sample_aspect_ratio = source->streams[stream_index]->sample_aspect_ratio;
    av_dict_copy(&output_stream->metadata, source->streams[stream_index]->metadata, 0);
  }
  av_dict_copy(&destination->metadata, source->metadata, 0);
  av_dict_set(&destination->metadata, "smartcut_title", "1", 0);
  if (!description.empty()) av_dict_set(&destination->metadata, "title", description.c_str(), 0);
  if (avio_open(&destination->pb, output.string().c_str(), AVIO_FLAG_WRITE) < 0 || avformat_write_header(destination, nullptr) < 0) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title header write failed\"}"; }
  frame->format = encoder->pix_fmt; frame->width = encoder->width; frame->height = encoder->height;
  if (av_frame_get_buffer(frame, 32) < 0) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title frame allocation failed\"}"; }
  scaler = sws_getContext(image_width, image_height, AV_PIX_FMT_BGRA, encoder->width, encoder->height, encoder->pix_fmt, SWS_BICUBIC, nullptr, nullptr, nullptr);
  if (!scaler) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title scaler setup failed\"}"; }
  std::vector<std::uint8_t> pixels(original.size());
  std::uint64_t packets_written = 0;
  auto drain = [&]() {
    while (avcodec_receive_packet(encoder, packet) >= 0) {
      av_packet_rescale_ts(packet, encoder->time_base, destination->streams[video_index]->time_base);
      packet->stream_index = video_index; packet->pos = -1; if (packet->duration <= 0) packet->duration = 1;
      if (av_interleaved_write_frame(destination, packet) < 0) { av_packet_unref(packet); return false; }
      ++packets_written; av_packet_unref(packet);
    }
    return true;
  };
  bool encoded = true;
  for (std::int64_t index = 0; index < frames && encoded; ++index) {
    const double seconds = static_cast<double>(index) / av_q2d(rate);
    double opacity = 1.0;
    if (fade_in_seconds > 0.0) opacity = std::min(opacity, seconds / fade_in_seconds);
    if (fade_out_seconds > 0.0) opacity = std::min(opacity, std::max(0.0, (duration_seconds - seconds) / fade_out_seconds));
    opacity = std::clamp(opacity, 0.0, 1.0);
    for (std::size_t offset = 0; offset < original.size(); offset += 4) {
      pixels[offset] = static_cast<std::uint8_t>(std::llround(original[offset] * opacity));
      pixels[offset + 1] = static_cast<std::uint8_t>(std::llround(original[offset + 1] * opacity));
      pixels[offset + 2] = static_cast<std::uint8_t>(std::llround(original[offset + 2] * opacity));
      pixels[offset + 3] = 255;
    }
    if (av_frame_make_writable(frame) < 0) { encoded = false; break; }
    const std::uint8_t* source_data[] = {pixels.data(), nullptr, nullptr, nullptr};
    const int source_lines[] = {image_width * 4, 0, 0, 0};
    sws_scale(scaler, source_data, source_lines, 0, image_height, frame->data, frame->linesize);
    frame->pts = index; frame->pict_type = AV_PICTURE_TYPE_NONE;
    encoded = avcodec_send_frame(encoder, frame) >= 0 && drain();
  }
  if (encoded) encoded = avcodec_send_frame(encoder, nullptr) >= 0 && drain();
  const bool trailer = encoded && av_write_trailer(destination) >= 0;
  if (!trailer) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED\",\"error\":\"title packet write failed\"}"; }
  if (destination->pb) avio_closep(&destination->pb);
  avformat_free_context(destination); destination = nullptr;
  AVFormatContext* validation = nullptr; AVCodecContext* decoder = nullptr; AVFrame* decoded = av_frame_alloc(); AVPacket* validation_packet = av_packet_alloc(); std::uint64_t decoded_frames = 0;
  if (validation_packet && decoded && avformat_open_input(&validation, output.string().c_str(), nullptr, nullptr) >= 0 && avformat_find_stream_info(validation, nullptr) >= 0) {
    const int validation_video = primary_video_stream(validation); const auto* decoder_codec = validation_video >= 0 ? ffmpeg::preferred_video_decoder(validation->streams[validation_video]->codecpar->codec_id) : nullptr;
    decoder = decoder_codec ? avcodec_alloc_context3(decoder_codec) : nullptr;
    if (decoder && avcodec_parameters_to_context(decoder, validation->streams[validation_video]->codecpar) >= 0 && avcodec_open2(decoder, decoder_codec, nullptr) >= 0) {
      while (av_read_frame(validation, validation_packet) >= 0) { if (validation_packet->stream_index == validation_video && avcodec_send_packet(decoder, validation_packet) >= 0) while (avcodec_receive_frame(decoder, decoded) >= 0) ++decoded_frames; av_packet_unref(validation_packet); }
      avcodec_send_packet(decoder, nullptr); while (avcodec_receive_frame(decoder, decoded) >= 0) ++decoded_frames;
    }
  }
  const bool valid = validation && validation->nb_streams == source_stream_count && decoded_frames == static_cast<std::uint64_t>(frames);
  av_packet_free(&validation_packet); av_frame_free(&decoded); avcodec_free_context(&decoder); avformat_close_input(&validation);
  if (!valid) { cleanup(); return "{\"schema\":\"title-media-v1\",\"status\":\"FAILED VALIDATION\",\"error\":\"title decode or stream layout mismatch\"}"; }
  completed = true; cleanup();
  std::ostringstream report;
  report << "{\"schema\":\"title-media-v1\",\"status\":\"PASSED\",\"frames\":" << frames
         << ",\"packets\":" << packets_written << ",\"streams\":" << source_stream_count
         << ",\"duration_seconds\":" << duration_seconds << ",\"silent_non_video_gap\":true}";
  return report.str();
#endif
}

}  // namespace smartcut::media
