#pragma once

#ifdef SMARTCUT_HAS_FFMPEG
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
}

#include <algorithm>
#include <charconv>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <string>
#include <string_view>
#include <vector>

namespace smartcut::media {

inline int pcm_depth(const AVCodecParameters* parameters) {
  if (!parameters) return 0;
  return parameters->bits_per_coded_sample > 0
             ? parameters->bits_per_coded_sample
             : parameters->bits_per_raw_sample;
}

// Matroska has no DVD LPCM mapping. Sixteen-bit DVD LPCM is a three-byte
// packet header followed by ordinary signed big-endian samples.
inline bool pcm_dvd_is_losslessly_normalizable(const AVCodecParameters* parameters) {
  return !parameters || parameters->codec_id != AV_CODEC_ID_PCM_DVD || pcm_depth(parameters) == 16;
}

inline int pcm_bluray_layout_code(const AVCodecParameters* parameters) {
  if (!parameters || parameters->ch_layout.order != AV_CHANNEL_ORDER_NATIVE) return 0;
  switch (parameters->ch_layout.u.mask) {
    case AV_CH_LAYOUT_MONO: return 1;
    case AV_CH_LAYOUT_STEREO: return 3;
    case AV_CH_LAYOUT_SURROUND: return 4;
    case AV_CH_LAYOUT_2_1: return 5;
    case AV_CH_LAYOUT_4POINT0: return 6;
    case AV_CH_LAYOUT_2_2: return 7;
    case AV_CH_LAYOUT_5POINT0: return 8;
    case AV_CH_LAYOUT_5POINT1: return 9;
    case AV_CH_LAYOUT_7POINT0: return 10;
    case AV_CH_LAYOUT_7POINT1: return 11;
    default: return 0;
  }
}

inline int pcm_bluray_frequency_code(int sample_rate) {
  switch (sample_rate) {
    case 48000: return 1;
    case 96000: return 4;
    case 192000: return 5;
    default: return 0;
  }
}

inline bool pcm_bluray_is_losslessly_normalizable(const AVCodecParameters* parameters) {
  if (!parameters || parameters->codec_id != AV_CODEC_ID_PCM_BLURAY) return true;
  const auto depth = pcm_depth(parameters);
  return (depth == 16 || depth == 24) && pcm_bluray_layout_code(parameters) != 0 &&
         pcm_bluray_frequency_code(parameters->sample_rate) != 0;
}

inline bool normalize_pcm_parameters_for_matroska(AVCodecParameters* parameters) {
  if (!parameters) return false;
  if (parameters->codec_id == AV_CODEC_ID_DOLBY_E || parameters->codec_id == AV_CODEC_ID_S302M) {
    parameters->codec_id = AV_CODEC_ID_PCM_S24LE;
    if (parameters->sample_rate <= 0) parameters->sample_rate = 48000;
    if (parameters->ch_layout.nb_channels <= 0) {
      const AVChannelLayout stereo = AV_CHANNEL_LAYOUT_STEREO;
      if (av_channel_layout_copy(&parameters->ch_layout, &stereo) < 0) return false;
    }
    parameters->bits_per_coded_sample = 24;
    parameters->bits_per_raw_sample = 24;
    parameters->block_align = parameters->ch_layout.nb_channels * 3;
  } else if (parameters->codec_id == AV_CODEC_ID_PCM_DVD) {
    if (!pcm_dvd_is_losslessly_normalizable(parameters)) return false;
    parameters->codec_id = AV_CODEC_ID_PCM_S16BE;
    parameters->bits_per_coded_sample = 16;
    parameters->bits_per_raw_sample = 16;
  } else if (parameters->codec_id == AV_CODEC_ID_PCM_BLURAY) {
    if (!pcm_bluray_is_losslessly_normalizable(parameters)) return false;
    const auto depth = pcm_depth(parameters);
    parameters->codec_id = depth == 16 ? AV_CODEC_ID_PCM_S16BE : AV_CODEC_ID_PCM_S24BE;
    parameters->bits_per_coded_sample = depth;
    parameters->bits_per_raw_sample = depth;
    parameters->block_align = parameters->ch_layout.nb_channels * depth / 8;
  }
  parameters->codec_tag = 0;
  return true;
}

inline constexpr const char* kPcmBlurayOrigin = "SMARTCUT_PCM_BLURAY";
inline constexpr const char* kPcmBlurayLayout = "SMARTCUT_PCM_BLURAY_LAYOUT";
inline constexpr const char* kBitstreamAudioOrigin = "SMARTCUT_BITSTREAM_AUDIO_ORIGIN";

inline bool parse_pcm_metadata_integer(const char* text, int* value) {
  if (!text || !value) return false;
  const std::string_view input(text);
  if (input.empty()) return false;
  int parsed{};
  const auto result = std::from_chars(input.data(), input.data() + input.size(), parsed);
  if (result.ec != std::errc{} || result.ptr != input.data() + input.size()) return false;
  *value = parsed;
  return true;
}

inline void mark_bitstream_audio_origin(AVStream* output, const AVStream* source) {
  if (!output || !source) return;
  if (source->codecpar->codec_id == AV_CODEC_ID_DOLBY_E)
    av_dict_set(&output->metadata, kBitstreamAudioOrigin, "dolby_e", 0);
  else if (source->codecpar->codec_id == AV_CODEC_ID_S302M)
    av_dict_set(&output->metadata, kBitstreamAudioOrigin, "s302m", 0);
}

inline bool bitstream_audio_origin(const AVStream* stream, AVCodecID* codec = nullptr) {
  if (!stream || stream->codecpar->codec_id != AV_CODEC_ID_PCM_S24LE) return false;
  const auto* marker = av_dict_get(stream->metadata, kBitstreamAudioOrigin, nullptr, 0);
  if (!marker) return false;
  AVCodecID original = AV_CODEC_ID_NONE;
  if (std::strcmp(marker->value, "dolby_e") == 0) original = AV_CODEC_ID_DOLBY_E;
  else if (std::strcmp(marker->value, "s302m") == 0) original = AV_CODEC_ID_S302M;
  if (original == AV_CODEC_ID_NONE) return false;
  if (codec) *codec = original;
  return true;
}

inline void mark_pcm_bluray_origin(AVStream* output, const AVStream* source) {
  if (!output || !source || source->codecpar->codec_id != AV_CODEC_ID_PCM_BLURAY) return;
  const auto depth = pcm_depth(source->codecpar);
  av_dict_set(&output->metadata, kPcmBlurayOrigin, std::to_string(depth).c_str(), 0);
  av_dict_set(&output->metadata, kPcmBlurayLayout,
              std::to_string(pcm_bluray_layout_code(source->codecpar)).c_str(), 0);
}

inline bool pcm_bluray_origin(const AVStream* stream, int* depth = nullptr,
                              int* layout_code = nullptr) {
  if (!stream) return false;
  const auto* marker = av_dict_get(stream->metadata, kPcmBlurayOrigin, nullptr, 0);
  const auto* layout = av_dict_get(stream->metadata, kPcmBlurayLayout, nullptr, 0);
  if (!marker || !layout ||
      (stream->codecpar->codec_id != AV_CODEC_ID_PCM_S16BE &&
       stream->codecpar->codec_id != AV_CODEC_ID_PCM_S24BE)) return false;
  int parsed_depth = 0;
  int parsed_layout = 0;
  if (!parse_pcm_metadata_integer(marker->value, &parsed_depth) ||
      !parse_pcm_metadata_integer(layout->value, &parsed_layout)) return false;
  if ((parsed_depth != 16 && parsed_depth != 24) || parsed_layout < 1 || parsed_layout > 11)
    return false;
  if (depth) *depth = parsed_depth;
  if (layout_code) *layout_code = parsed_layout;
  return true;
}

// Maps each coded Blu-ray channel slot to the canonical FFmpeg channel slot.
// A value of -1 is the padding slot used by layouts with an odd channel count.
inline std::vector<int> pcm_bluray_coded_to_canonical(int layout_code) {
  switch (layout_code) {
    case 1: return {0, -1};
    case 3: return {0, 1};
    case 4: return {0, 1, 2, -1};
    case 5: return {0, 1, 2, -1};
    case 6: return {0, 1, 2, 3};
    case 7: return {0, 1, 2, 3};
    case 8: return {0, 1, 2, 3, 4, -1};
    case 9: return {0, 1, 2, 4, 5, 3};
    case 10: return {0, 1, 2, 5, 3, 4, 6, -1};
    case 11: return {0, 1, 2, 6, 4, 5, 7, 3};
    default: return {};
  }
}

inline std::vector<int> pcm_bluray_coded_to_canonical(const AVCodecParameters* parameters) {
  return pcm_bluray_coded_to_canonical(pcm_bluray_layout_code(parameters));
}

inline bool set_pcm_bluray_channel_layout(AVCodecParameters* parameters, int layout_code) {
  if (!parameters) return false;
  av_channel_layout_uninit(&parameters->ch_layout);
  switch (layout_code) {
    case 1: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_MONO; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 3: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_STEREO; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 4: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_SURROUND; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 5: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_2_1; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 6: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_4POINT0; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 7: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_2_2; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 8: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_5POINT0; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 9: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_5POINT1; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 10: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_7POINT0; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    case 11: { const AVChannelLayout layout = AV_CHANNEL_LAYOUT_7POINT1; return av_channel_layout_copy(&parameters->ch_layout, &layout) == 0; }
    default: return false;
  }
}

inline bool replace_packet_payload(AVPacket* packet, const std::vector<std::uint8_t>& payload) {
  if (!packet || payload.empty()) return false;
  auto* replacement = av_packet_alloc();
  if (!replacement || av_new_packet(replacement, static_cast<int>(payload.size())) < 0 ||
      av_packet_copy_props(replacement, packet) < 0) {
    av_packet_free(&replacement);
    return false;
  }
  std::memcpy(replacement->data, payload.data(), payload.size());
  replacement->stream_index = packet->stream_index;
  replacement->pts = packet->pts;
  replacement->dts = packet->dts;
  replacement->duration = packet->duration;
  replacement->pos = packet->pos;
  replacement->flags = packet->flags;
  av_packet_unref(packet);
  av_packet_move_ref(packet, replacement);
  av_packet_free(&replacement);
  return true;
}

inline std::vector<std::uint8_t> normalized_pcm_bluray_payload(
    const AVPacket* packet, const AVCodecParameters* source) {
  if (!packet || !source || packet->size < 4 ||
      source->codec_id != AV_CODEC_ID_PCM_BLURAY ||
      !pcm_bluray_is_losslessly_normalizable(source)) return {};
  const auto depth = pcm_depth(source);
  const auto bytes_per_sample = static_cast<std::size_t>(depth / 8);
  const auto mapping = pcm_bluray_coded_to_canonical(source);
  const auto channels = static_cast<std::size_t>(source->ch_layout.nb_channels);
  const auto coded_channels = mapping.size();
  const auto payload_size = static_cast<std::size_t>(packet->size - 4);
  const auto coded_sample_size = coded_channels * bytes_per_sample;
  const auto expected_layout = pcm_bluray_layout_code(source);
  const auto expected_frequency = pcm_bluray_frequency_code(source->sample_rate);
  const auto expected_quantization = depth == 16 ? 1 : 3;
  if (mapping.empty() || channels == 0 || coded_sample_size == 0 ||
      payload_size % coded_sample_size != 0 ||
      (packet->data[2] >> 4) != expected_layout ||
      (packet->data[2] & 0x0f) != expected_frequency ||
      (packet->data[3] >> 6) != expected_quantization) return {};
  const auto samples = payload_size / coded_sample_size;
  std::vector<std::uint8_t> normalized(samples * channels * bytes_per_sample);
  const auto* input = packet->data + 4;
  for (std::size_t sample = 0; sample < samples; ++sample) {
    for (std::size_t coded = 0; coded < coded_channels; ++coded) {
      const auto canonical = mapping[coded];
      if (canonical < 0) continue;
      std::memcpy(normalized.data() + (sample * channels + static_cast<std::size_t>(canonical)) * bytes_per_sample,
                  input + (sample * coded_channels + coded) * bytes_per_sample,
                  bytes_per_sample);
    }
  }
  return normalized;
}

inline bool normalize_pcm_packet_for_matroska(AVPacket* packet,
                                               const AVCodecParameters* source) {
  if (!packet || !source) return false;
  if (source->codec_id == AV_CODEC_ID_PCM_DVD) {
    if (!pcm_dvd_is_losslessly_normalizable(source) || packet->size < 3) return false;
    std::vector<std::uint8_t> payload(packet->data + 3, packet->data + packet->size);
    return replace_packet_payload(packet, payload);
  }
  if (source->codec_id == AV_CODEC_ID_PCM_BLURAY) {
    const auto payload = normalized_pcm_bluray_payload(packet, source);
    return replace_packet_payload(packet, payload);
  }
  return true;
}

inline bool restore_pcm_bluray_packet(AVPacket* packet, const AVCodecParameters* normalized,
                                      int layout_code = 0) {
  if (!packet || !normalized || packet->size <= 0 ||
      (normalized->codec_id != AV_CODEC_ID_PCM_S16BE &&
       normalized->codec_id != AV_CODEC_ID_PCM_S24BE)) return false;
  const auto depth = normalized->codec_id == AV_CODEC_ID_PCM_S16BE ? 16 : 24;
  const auto bytes_per_sample = static_cast<std::size_t>(depth / 8);
  if (layout_code == 0) layout_code = pcm_bluray_layout_code(normalized);
  const auto mapping = pcm_bluray_coded_to_canonical(layout_code);
  const auto channels = static_cast<std::size_t>(normalized->ch_layout.nb_channels);
  const auto coded_channels = mapping.size();
  const auto standard_sample_size = channels * bytes_per_sample;
  if (mapping.empty() || channels == 0 || standard_sample_size == 0 ||
      static_cast<std::size_t>(packet->size) % standard_sample_size != 0) return false;
  const auto samples = static_cast<std::size_t>(packet->size) / standard_sample_size;
  const auto coded_payload_size = samples * coded_channels * bytes_per_sample;
  if (coded_payload_size > 0xffffU) return false;
  std::vector<std::uint8_t> restored(coded_payload_size + 4U, 0);
  restored[0] = static_cast<std::uint8_t>((coded_payload_size >> 8) & 0xffU);
  restored[1] = static_cast<std::uint8_t>(coded_payload_size & 0xffU);
  restored[2] = static_cast<std::uint8_t>((layout_code << 4) |
                                          pcm_bluray_frequency_code(normalized->sample_rate));
  restored[3] = static_cast<std::uint8_t>((depth == 16 ? 1 : 3) << 6);
  auto* output = restored.data() + 4;
  for (std::size_t sample = 0; sample < samples; ++sample) {
    for (std::size_t coded = 0; coded < coded_channels; ++coded) {
      const auto canonical = mapping[coded];
      if (canonical < 0) continue;
      std::memcpy(output + (sample * coded_channels + coded) * bytes_per_sample,
                  packet->data + (sample * channels + static_cast<std::size_t>(canonical)) * bytes_per_sample,
                  bytes_per_sample);
    }
  }
  return replace_packet_payload(packet, restored);
}

inline bool matroska_codec_equivalent(const AVCodecParameters* source,
                                      const AVCodecParameters* rendered) {
  if (!source || !rendered || source->codec_type != rendered->codec_type) return false;
  if (source->codec_id == rendered->codec_id) return true;
  if (source->codec_id == AV_CODEC_ID_PCM_DVD)
    return rendered->codec_id == AV_CODEC_ID_PCM_S16BE && pcm_dvd_is_losslessly_normalizable(source);
  if (source->codec_id == AV_CODEC_ID_PCM_BLURAY) {
    const auto expected = pcm_depth(source) == 16 ? AV_CODEC_ID_PCM_S16BE : AV_CODEC_ID_PCM_S24BE;
    return rendered->codec_id == expected && pcm_bluray_is_losslessly_normalizable(source);
  }
  if (source->codec_id == AV_CODEC_ID_DOLBY_E || source->codec_id == AV_CODEC_ID_S302M)
    return rendered->codec_id == AV_CODEC_ID_PCM_S24LE;
  return false;
}

inline std::vector<std::uint8_t> matroska_normalized_payload(
    const AVPacket* packet, const AVCodecParameters* source) {
  if (!packet || packet->size <= 0) return {};
  if (source && source->codec_id == AV_CODEC_ID_PCM_DVD) {
    if (packet->size < 3) return {};
    return {packet->data + 3, packet->data + packet->size};
  }
  if (source && source->codec_id == AV_CODEC_ID_PCM_BLURAY)
    return normalized_pcm_bluray_payload(packet, source);
  return {packet->data, packet->data + packet->size};
}

}  // namespace smartcut::media
#endif
