#include "media/audio_transform.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <limits>
#include <sstream>
#include <system_error>
#include <vector>

#ifdef SMARTCUT_HAS_FFMPEG
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/audio_fifo.h>
#include <libavutil/samplefmt.h>
#include <libswresample/swresample.h>
}
#include "media/ffmpeg_raii.hpp"
#endif

namespace smartcut::media {
namespace {

std::string failure(const std::string& reason) {
  return "{\"schema\":\"audio-transform-v1\",\"status\":\"FAILED\",\"reason\":\"" + reason + "\"}";
}

void remove_file(const std::string& path) {
  std::error_code ignored;
  std::filesystem::remove(path, ignored);
}

#ifdef SMARTCUT_HAS_FFMPEG
struct AudioFifoDeleter { void operator()(AVAudioFifo* fifo) const { av_audio_fifo_free(fifo); } };
using AudioFifo = std::unique_ptr<AVAudioFifo, AudioFifoDeleter>;

const AVCodec* encoder_for(AudioOutputCodec codec, AVCodecID source_codec) {
  if (codec == AudioOutputCodec::source) return avcodec_find_encoder(source_codec);
  if (codec == AudioOutputCodec::aac) return avcodec_find_encoder(AV_CODEC_ID_AAC);
  if (codec == AudioOutputCodec::ac3) return avcodec_find_encoder(AV_CODEC_ID_AC3);
  if (codec == AudioOutputCodec::mp2) return avcodec_find_encoder(AV_CODEC_ID_MP2);
  if (codec == AudioOutputCodec::pcm_s16le) return avcodec_find_encoder(AV_CODEC_ID_PCM_S16LE);
  if (codec == AudioOutputCodec::opus) {
    if (const auto* encoder = avcodec_find_encoder_by_name("libopus")) return encoder;
    return avcodec_find_encoder(AV_CODEC_ID_OPUS);
  }
  return nullptr;
}

AVSampleFormat choose_sample_format(const AVCodec* codec) {
  const void* values = nullptr;
  int count = 0;
  if (avcodec_get_supported_config(nullptr, codec, AV_CODEC_CONFIG_SAMPLE_FORMAT, 0, &values, &count) < 0 || !values || count == 0)
    return AV_SAMPLE_FMT_FLTP;
  const auto* formats = static_cast<const AVSampleFormat*>(values);
  for (const auto preferred : {AV_SAMPLE_FMT_FLTP, AV_SAMPLE_FMT_FLT, AV_SAMPLE_FMT_S16P, AV_SAMPLE_FMT_S16})
    for (int index = 0; index < count; ++index) if (formats[index] == preferred) return preferred;
  return formats[0];
}

int choose_sample_rate(const AVCodec* codec, int requested) {
  const int wanted = requested > 0 ? requested : 48000;
  const void* values = nullptr;
  int count = 0;
  if (avcodec_get_supported_config(nullptr, codec, AV_CODEC_CONFIG_SAMPLE_RATE, 0, &values, &count) < 0 || !values || count == 0)
    return wanted;
  const auto* rates = static_cast<const int*>(values);
  int selected = rates[0];
  for (int index = 0; index < count; ++index)
    if (std::abs(rates[index] - wanted) < std::abs(selected - wanted)) selected = rates[index];
  return selected;
}

void copy_stream_attributes(AVStream* output, const AVStream* input) {
  output->disposition = input->disposition;
  output->avg_frame_rate = input->avg_frame_rate;
  output->sample_aspect_ratio = input->sample_aspect_ratio;
  av_dict_copy(&output->metadata, input->metadata, 0);
}

template <typename Sample>
void scale_integer_samples(std::uint8_t* data, int count, double gain) {
  auto* samples = reinterpret_cast<Sample*>(data);
  const auto minimum = static_cast<double>(std::numeric_limits<Sample>::min());
  const auto maximum = static_cast<double>(std::numeric_limits<Sample>::max());
  for (int index = 0; index < count; ++index)
    samples[index] = static_cast<Sample>(std::clamp(std::nearbyint(static_cast<double>(samples[index]) * gain), minimum, maximum));
}

void apply_gain(AVFrame* frame, double gain) {
  if (std::abs(gain - 1.0) < 0.000001) return;
  const auto format = static_cast<AVSampleFormat>(frame->format);
  const bool planar = av_sample_fmt_is_planar(format) != 0;
  const int channels = frame->ch_layout.nb_channels;
  const int planes = planar ? channels : 1;
  const int samples_per_plane = planar ? frame->nb_samples : frame->nb_samples * channels;
  const auto packed = av_get_packed_sample_fmt(format);
  for (int plane = 0; plane < planes; ++plane) {
    auto* data = frame->extended_data[plane];
    if (packed == AV_SAMPLE_FMT_U8) {
      for (int index = 0; index < samples_per_plane; ++index)
        data[index] = static_cast<std::uint8_t>(std::clamp(std::nearbyint((static_cast<double>(data[index]) - 128.0) * gain + 128.0), 0.0, 255.0));
    } else if (packed == AV_SAMPLE_FMT_S16) scale_integer_samples<std::int16_t>(data, samples_per_plane, gain);
    else if (packed == AV_SAMPLE_FMT_S32) scale_integer_samples<std::int32_t>(data, samples_per_plane, gain);
    else if (packed == AV_SAMPLE_FMT_S64) scale_integer_samples<std::int64_t>(data, samples_per_plane, gain);
    else if (packed == AV_SAMPLE_FMT_FLT) { auto* samples = reinterpret_cast<float*>(data); for (int index = 0; index < samples_per_plane; ++index) samples[index] = static_cast<float>(std::clamp(samples[index] * gain, -1.0, 1.0)); }
    else if (packed == AV_SAMPLE_FMT_DBL) { auto* samples = reinterpret_cast<double*>(data); for (int index = 0; index < samples_per_plane; ++index) samples[index] = std::clamp(samples[index] * gain, -1.0, 1.0); }
  }
}

void copy_chapters(AVFormatContext* output, const AVFormatContext* input) {
  for (unsigned index = 0; index < input->nb_chapters; ++index) {
    const auto* chapter = input->chapters[index];
    auto* copied = static_cast<AVChapter*>(av_mallocz(sizeof(AVChapter)));
    if (!copied) continue;
    copied->id = chapter->id; copied->time_base = chapter->time_base; copied->start = chapter->start; copied->end = chapter->end;
    av_dict_copy(&copied->metadata, chapter->metadata, 0);
    auto** chapters = static_cast<AVChapter**>(av_realloc_array(output->chapters, output->nb_chapters + 1, sizeof(*output->chapters)));
    if (!chapters) { av_dict_free(&copied->metadata); av_free(copied); continue; }
    output->chapters = chapters; output->chapters[output->nb_chapters++] = copied;
  }
}
#endif

}  // namespace

const char* audio_output_codec_name(AudioOutputCodec codec) {
  switch (codec) {
    case AudioOutputCodec::copy: return "copy";
    case AudioOutputCodec::source: return "source";
    case AudioOutputCodec::aac: return "aac";
    case AudioOutputCodec::ac3: return "ac3";
    case AudioOutputCodec::mp2: return "mp2";
    case AudioOutputCodec::opus: return "opus";
    case AudioOutputCodec::pcm_s16le: return "pcm_s16le";
  }
  return "copy";
}

AudioOutputCodec audio_output_codec_from_name(const std::string& name) {
  if (name == "source") return AudioOutputCodec::source;
  if (name == "aac") return AudioOutputCodec::aac;
  if (name == "ac3") return AudioOutputCodec::ac3;
  if (name == "mp2") return AudioOutputCodec::mp2;
  if (name == "opus") return AudioOutputCodec::opus;
  if (name == "pcm" || name == "pcm_s16le" || name == "lpcm") return AudioOutputCodec::pcm_s16le;
  return AudioOutputCodec::copy;
}

bool audio_transform_requested(const AudioTransformOptions& options) {
  return options.codec != AudioOutputCodec::copy || std::abs(options.gain_db) > 0.000001 || options.sample_rate > 0 || options.channels > 0;
}

std::string transform_primary_audio_mkv(const std::string& input, const std::string& output,
                                        const AudioTransformOptions& options) {
#ifndef SMARTCUT_HAS_FFMPEG
  (void)input; (void)output; (void)options;
  return failure("FFmpeg support is not enabled");
#else
  if (options.codec == AudioOutputCodec::copy) return failure("audio gain or format changes require an encoding codec");
  if (options.bitrate < 16000 || options.bitrate > 1536000 || options.gain_db < -60.0 || options.gain_db > 24.0)
    return failure("audio profile values are outside supported limits");
  remove_file(output);
  AVFormatContext* source = nullptr;
  AVFormatContext* destination = nullptr;
  AVAudioFifo* raw_fifo = nullptr;
  auto input_packet = ffmpeg::make_packet();
  auto encoded_packet = ffmpeg::make_packet();
  auto decoded = ffmpeg::make_frame();
  bool completed = false;
  auto cleanup = [&] {
    if (raw_fifo) av_audio_fifo_free(raw_fifo);
    if (destination && destination->pb) avio_closep(&destination->pb);
    avformat_free_context(destination);
    avformat_close_input(&source);
    if (!completed) remove_file(output);
  };
  if (!input_packet || !encoded_packet || !decoded || avformat_open_input(&source, input.c_str(), nullptr, nullptr) < 0 ||
      avformat_find_stream_info(source, nullptr) < 0) { cleanup(); return failure("input open failed"); }
  int audio_index = -1;
  for (unsigned index = 0; index < source->nb_streams; ++index)
    if (source->streams[index]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) { audio_index = static_cast<int>(index); break; }
  if (audio_index < 0) { cleanup(); return failure("source has no audio stream"); }

  const auto* input_audio = source->streams[audio_index];
  const auto* decoder_codec = avcodec_find_decoder(input_audio->codecpar->codec_id);
  const auto* encoder_codec = encoder_for(options.codec, input_audio->codecpar->codec_id);
  auto decoder = ffmpeg::make_codec_context(decoder_codec);
  auto encoder = ffmpeg::make_codec_context(encoder_codec);
  if (!decoder_codec || !encoder_codec || !decoder || !encoder ||
      avcodec_parameters_to_context(decoder.get(), input_audio->codecpar) < 0 || avcodec_open2(decoder.get(), decoder_codec, nullptr) < 0) {
    cleanup(); return failure("audio decoder or encoder is unavailable");
  }
  if (decoder->ch_layout.nb_channels <= 0)
    av_channel_layout_default(&decoder->ch_layout, input_audio->codecpar->ch_layout.nb_channels > 0 ? input_audio->codecpar->ch_layout.nb_channels : 2);
  const int output_channels = options.channels > 0 ? options.channels : decoder->ch_layout.nb_channels;
  av_channel_layout_default(&encoder->ch_layout, std::clamp(output_channels, 1, 8));
  encoder->sample_rate = choose_sample_rate(encoder_codec, options.sample_rate > 0 ? options.sample_rate : decoder->sample_rate);
  encoder->sample_fmt = choose_sample_format(encoder_codec);
  encoder->bit_rate = options.bitrate;
  encoder->time_base = AVRational{1, encoder->sample_rate};
  if (options.codec == AudioOutputCodec::opus) encoder->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
  if (avformat_alloc_output_context2(&destination, nullptr, "matroska", output.c_str()) < 0 || !destination) {
    cleanup(); return failure("output context failed");
  }
  if (destination->oformat->flags & AVFMT_GLOBALHEADER) encoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  if (avcodec_open2(encoder.get(), encoder_codec, nullptr) < 0) { cleanup(); return failure("audio encoder initialization failed"); }

  std::vector<int> stream_map(source->nb_streams, -1);
  std::vector<AVCodecID> copied_codecs(source->nb_streams, AV_CODEC_ID_NONE);
  AVStream* output_audio = nullptr;
  for (unsigned index = 0; index < source->nb_streams; ++index) {
    const auto* input_stream = source->streams[index];
    auto* output_stream = avformat_new_stream(destination, nullptr);
    if (!output_stream) { cleanup(); return failure("output stream allocation failed"); }
    stream_map[index] = output_stream->index;
    if (static_cast<int>(index) == audio_index) {
      if (avcodec_parameters_from_context(output_stream->codecpar, encoder.get()) < 0) { cleanup(); return failure("encoded audio stream setup failed"); }
      output_stream->time_base = encoder->time_base;
      output_audio = output_stream;
    } else {
      if (avcodec_parameters_copy(output_stream->codecpar, input_stream->codecpar) < 0) { cleanup(); return failure("copied stream setup failed"); }
      output_stream->codecpar->codec_tag = 0;
      output_stream->time_base = input_stream->time_base;
      copied_codecs[index] = input_stream->codecpar->codec_id;
    }
    copy_stream_attributes(output_stream, input_stream);
    if (static_cast<int>(index) == audio_index) av_dict_set(&output_stream->metadata, "ENCODER", nullptr, 0);
  }
  av_dict_copy(&destination->metadata, source->metadata, 0);
  copy_chapters(destination, source);
  if (!(destination->oformat->flags & AVFMT_NOFILE) && avio_open(&destination->pb, output.c_str(), AVIO_FLAG_WRITE) < 0) { cleanup(); return failure("output open failed"); }
  if (avformat_write_header(destination, nullptr) < 0) { cleanup(); return failure("output header write failed"); }

  SwrContext* raw_resampler = nullptr;
  if (swr_alloc_set_opts2(&raw_resampler, &encoder->ch_layout, encoder->sample_fmt, encoder->sample_rate,
                          &decoder->ch_layout, decoder->sample_fmt, decoder->sample_rate, 0, nullptr) < 0) {
    cleanup(); return failure("audio resampler allocation failed");
  }
  ffmpeg::Swr resampler(raw_resampler);
  if (swr_init(resampler.get()) < 0) { cleanup(); return failure("audio resampler initialization failed"); }
  raw_fifo = av_audio_fifo_alloc(encoder->sample_fmt, encoder->ch_layout.nb_channels, std::max(encoder->frame_size, 1024));
  if (!raw_fifo) { cleanup(); return failure("audio FIFO allocation failed"); }

  std::int64_t next_pts = 0;
  std::size_t encoded_packets = 0;
  std::vector<std::size_t> copied_packets(source->nb_streams, 0);
  bool failed = false;
  const double gain = std::pow(10.0, options.gain_db / 20.0);
  auto write_encoded = [&] {
    while (avcodec_receive_packet(encoder.get(), encoded_packet.get()) >= 0) {
      av_packet_rescale_ts(encoded_packet.get(), encoder->time_base, output_audio->time_base);
      encoded_packet->stream_index = output_audio->index; encoded_packet->pos = -1;
      if (av_interleaved_write_frame(destination, encoded_packet.get()) < 0) failed = true;
      else ++encoded_packets;
      av_packet_unref(encoded_packet.get());
    }
  };
  auto encode_fifo = [&](bool flush) {
    const int nominal = encoder->frame_size > 0 ? encoder->frame_size : 1024;
    while (!failed && av_audio_fifo_size(raw_fifo) >= nominal) {
      auto frame = ffmpeg::make_frame(); if (!frame) { failed = true; break; }
      frame->format = encoder->sample_fmt; frame->sample_rate = encoder->sample_rate; av_channel_layout_copy(&frame->ch_layout, &encoder->ch_layout); frame->nb_samples = nominal; frame->pts = next_pts;
      if (av_frame_get_buffer(frame.get(), 0) < 0 || av_audio_fifo_read(raw_fifo, reinterpret_cast<void**>(frame->data), nominal) < nominal) { failed = true; break; }
      next_pts += nominal; if (avcodec_send_frame(encoder.get(), frame.get()) < 0) { failed = true; break; } write_encoded();
    }
    if (flush && !failed && av_audio_fifo_size(raw_fifo) > 0) {
      const int available = av_audio_fifo_size(raw_fifo);
      const bool short_frame = (encoder_codec->capabilities & (AV_CODEC_CAP_SMALL_LAST_FRAME | AV_CODEC_CAP_VARIABLE_FRAME_SIZE)) != 0;
      const int submitted = short_frame ? available : nominal;
      auto frame = ffmpeg::make_frame(); if (!frame) { failed = true; return; }
      frame->format = encoder->sample_fmt; frame->sample_rate = encoder->sample_rate; av_channel_layout_copy(&frame->ch_layout, &encoder->ch_layout); frame->nb_samples = submitted; frame->pts = next_pts;
      if (av_frame_get_buffer(frame.get(), 0) < 0 || av_audio_fifo_read(raw_fifo, reinterpret_cast<void**>(frame->data), available) < available) { failed = true; return; }
      if (submitted > available) av_samples_set_silence(frame->data, available, submitted - available, encoder->ch_layout.nb_channels, encoder->sample_fmt);
      next_pts += available; if (avcodec_send_frame(encoder.get(), frame.get()) < 0) { failed = true; return; } write_encoded();
    }
  };
  auto convert_decoded = [&] {
    while (!failed && avcodec_receive_frame(decoder.get(), decoded.get()) >= 0) {
      const int capacity = static_cast<int>(av_rescale_rnd(swr_get_delay(resampler.get(), decoder->sample_rate) + decoded->nb_samples,
                                                           encoder->sample_rate, decoder->sample_rate, AV_ROUND_UP));
      auto converted = ffmpeg::make_frame(); if (!converted) { failed = true; break; }
      converted->format = encoder->sample_fmt; converted->sample_rate = encoder->sample_rate; av_channel_layout_copy(&converted->ch_layout, &encoder->ch_layout); converted->nb_samples = capacity;
      if (av_frame_get_buffer(converted.get(), 0) < 0) { failed = true; break; }
      const int samples = swr_convert(resampler.get(), converted->data, capacity, const_cast<const std::uint8_t**>(decoded->extended_data), decoded->nb_samples);
      av_frame_unref(decoded.get());
      if (samples < 0) { failed = true; break; }
      converted->nb_samples = samples; apply_gain(converted.get(), gain);
      if (av_audio_fifo_realloc(raw_fifo, av_audio_fifo_size(raw_fifo) + samples) < 0 || av_audio_fifo_write(raw_fifo, reinterpret_cast<void**>(converted->data), samples) < samples) { failed = true; break; }
      encode_fifo(false);
    }
  };

  while (!failed && av_read_frame(source, input_packet.get()) >= 0) {
    const int source_index = input_packet->stream_index;
    if (source_index == audio_index) {
      if (avcodec_send_packet(decoder.get(), input_packet.get()) < 0) failed = true;
      else convert_decoded();
    } else {
      auto* output_stream = destination->streams[stream_map[static_cast<std::size_t>(source_index)]];
      av_packet_rescale_ts(input_packet.get(), source->streams[source_index]->time_base, output_stream->time_base);
      input_packet->stream_index = output_stream->index; input_packet->pos = -1;
      if (av_interleaved_write_frame(destination, input_packet.get()) < 0) failed = true;
      else ++copied_packets[static_cast<std::size_t>(source_index)];
    }
    av_packet_unref(input_packet.get());
  }
  if (!failed) { avcodec_send_packet(decoder.get(), nullptr); convert_decoded(); }
  if (!failed) {
    const int delayed = static_cast<int>(swr_get_delay(resampler.get(), decoder->sample_rate));
    if (delayed > 0) {
      const int capacity = static_cast<int>(av_rescale_rnd(delayed, encoder->sample_rate, decoder->sample_rate, AV_ROUND_UP));
      auto tail = ffmpeg::make_frame();
      if (!tail) failed = true;
      else { tail->format = encoder->sample_fmt; tail->sample_rate = encoder->sample_rate; av_channel_layout_copy(&tail->ch_layout, &encoder->ch_layout); tail->nb_samples = capacity; if (av_frame_get_buffer(tail.get(), 0) < 0) failed = true; else { const int samples = swr_convert(resampler.get(), tail->data, capacity, nullptr, 0); if (samples < 0) failed = true; else { tail->nb_samples = samples; apply_gain(tail.get(), gain); if (av_audio_fifo_realloc(raw_fifo, av_audio_fifo_size(raw_fifo) + samples) < 0 || av_audio_fifo_write(raw_fifo, reinterpret_cast<void**>(tail->data), samples) < samples) failed = true; } } }
    }
  }
  if (!failed) encode_fifo(true);
  if (!failed) { avcodec_send_frame(encoder.get(), nullptr); write_encoded(); }
  if (!failed && av_write_trailer(destination) < 0) failed = true;
  if (failed || encoded_packets == 0) { cleanup(); return failure("audio decode, encode, or mux failed"); }

  if (destination && destination->pb) avio_closep(&destination->pb);
  avformat_free_context(destination); destination = nullptr; avformat_close_input(&source);
  AVFormatContext* validation = nullptr;
  if (avformat_open_input(&validation, output.c_str(), nullptr, nullptr) < 0 || avformat_find_stream_info(validation, nullptr) < 0) {
    avformat_close_input(&validation); cleanup(); return failure("output validation open failed");
  }
  bool valid = validation->nb_streams == stream_map.size() && validation->streams[static_cast<unsigned>(audio_index)]->codecpar->codec_id == encoder->codec_id;
  for (std::size_t index = 0; valid && index < copied_codecs.size(); ++index)
    if (copied_codecs[index] != AV_CODEC_ID_NONE) valid = validation->streams[index]->codecpar->codec_id == copied_codecs[index];
  const auto output_codec = validation->streams[static_cast<unsigned>(audio_index)]->codecpar->codec_id;
  const int output_rate = validation->streams[static_cast<unsigned>(audio_index)]->codecpar->sample_rate;
  const int output_channel_count = validation->streams[static_cast<unsigned>(audio_index)]->codecpar->ch_layout.nb_channels;
  avformat_close_input(&validation);
  if (!valid) { cleanup(); return failure("output codec validation failed"); }
  completed = true;
  std::ostringstream report;
  report << "{\"schema\":\"audio-transform-v1\",\"status\":\"PASSED\",\"primary_stream\":" << audio_index
         << ",\"codec\":\"" << avcodec_get_name(output_codec) << "\",\"gain_db\":" << options.gain_db
         << ",\"bitrate\":" << options.bitrate << ",\"sample_rate\":" << output_rate
         << ",\"channels\":" << output_channel_count << ",\"encoded_packets\":" << encoded_packets
         << ",\"video_policy\":\"packet-copy\",\"other_audio_policy\":\"packet-copy\"}";
  cleanup();
  return report.str();
#endif
}

}  // namespace smartcut::media
