#include "media/media_joiner.hpp"

#include "media/container_mux.hpp"

#include <algorithm>
#include <chrono>
#include <cstdint>
#include <cstring>
#include <limits>
#include <sstream>
#include <string_view>
#include <system_error>
#include <vector>

#include "media/ffmpeg_raii.hpp"

namespace smartcut::media {
namespace {
#ifdef SMARTCUT_HAS_FFMPEG
struct InputTiming {
  std::int64_t first_us{};
  std::int64_t end_us{};
  std::int64_t video_first_us{AV_NOPTS_VALUE};
  std::int64_t video_end_us{AV_NOPTS_VALUE};
  std::uint64_t video_packets{};
  std::vector<std::int64_t> first_dts_us;
};
struct ChapterCopy { std::int64_t start_us{}; std::int64_t end_us{}; std::string title; };

bool compatible_streams(const AVCodecParameters* first, const AVCodecParameters* next) {
  if (first->codec_type != next->codec_type || first->codec_id != next->codec_id) return false;
  if (first->codec_type == AVMEDIA_TYPE_VIDEO)
    return first->width == next->width && first->height == next->height && first->format == next->format &&
           first->profile == next->profile;
  if (first->codec_type == AVMEDIA_TYPE_AUDIO &&
      (first->sample_rate != next->sample_rate ||
       first->ch_layout.nb_channels != next->ch_layout.nb_channels)) return false;
  if (first->extradata_size != next->extradata_size) return false;
  return first->extradata_size == 0 ||
         std::memcmp(first->extradata, next->extradata,
                     static_cast<std::size_t>(first->extradata_size)) == 0;
}

InputTiming scan_timing(AVFormatContext* input) {
  if (av_seek_frame(input, -1, 0, AVSEEK_FLAG_BACKWARD) >= 0) avformat_flush(input);
  std::int64_t first = std::numeric_limits<std::int64_t>::max();
  std::int64_t end = std::numeric_limits<std::int64_t>::min();
  std::int64_t video_first = std::numeric_limits<std::int64_t>::max();
  std::int64_t video_end = std::numeric_limits<std::int64_t>::min();
  std::uint64_t video_packets = 0;
  std::vector<std::int64_t> first_dts(input->nb_streams, AV_NOPTS_VALUE);
  for (unsigned index = 0; index < input->nb_streams; ++index) {
    const auto* stream = input->streams[index];
    if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO &&
        !(stream->disposition & AV_DISPOSITION_ATTACHED_PIC) && stream->start_time != AV_NOPTS_VALUE)
      video_first = std::min(video_first,
                             av_rescale_q(stream->start_time, stream->time_base, AV_TIME_BASE_Q));
  }
  auto packet = ffmpeg::make_packet();
  while (packet && av_read_frame(input, packet.get()) >= 0) {
    const auto* stream = input->streams[packet->stream_index];
    if (packet->dts != AV_NOPTS_VALUE && first_dts[static_cast<std::size_t>(packet->stream_index)] == AV_NOPTS_VALUE)
      first_dts[static_cast<std::size_t>(packet->stream_index)] =
          av_rescale_q(packet->dts, stream->time_base, AV_TIME_BASE_Q);
    const auto timestamp = packet->dts != AV_NOPTS_VALUE ? packet->dts : packet->pts;
    const auto presentation = packet->pts != AV_NOPTS_VALUE ? packet->pts : timestamp;
    if (timestamp != AV_NOPTS_VALUE) first = std::min(first, av_rescale_q(timestamp, stream->time_base, AV_TIME_BASE_Q));
    if (presentation != AV_NOPTS_VALUE) {
      const auto packet_end = presentation + std::max<std::int64_t>(packet->duration, 0);
      end = std::max(end, av_rescale_q(packet_end, stream->time_base, AV_TIME_BASE_Q));
      if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO &&
          !(stream->disposition & AV_DISPOSITION_ATTACHED_PIC)) {
        ++video_packets;
        video_first = std::min(video_first,
                               av_rescale_q(presentation, stream->time_base, AV_TIME_BASE_Q));
        video_end = std::max(video_end,
                             av_rescale_q(packet_end, stream->time_base, AV_TIME_BASE_Q));
      }
    }
    av_packet_unref(packet.get());
  }
  if (first == std::numeric_limits<std::int64_t>::max()) first = 0;
  if (end <= first) end = first + std::max<std::int64_t>(input->duration, 1);
  if (video_first == std::numeric_limits<std::int64_t>::max()) video_first = AV_NOPTS_VALUE;
  if (video_end == std::numeric_limits<std::int64_t>::min()) video_end = AV_NOPTS_VALUE;
  for (unsigned index = 0; index < input->nb_streams; ++index) {
    const auto* stream = input->streams[index];
    if (stream->codecpar->codec_type != AVMEDIA_TYPE_VIDEO ||
        (stream->disposition & AV_DISPOSITION_ATTACHED_PIC)) continue;
    const auto* tagged = av_dict_get(stream->metadata, "SMARTCUT_FRAME_COUNT", nullptr, 0);
    if (tagged && tagged->value) {
      try {
        const auto parsed = std::stoull(tagged->value);
        if (parsed > 0) video_packets = parsed;
      } catch (...) {}
    }
    break;
  }
  return {first, end, video_first, video_end, video_packets, std::move(first_dts)};
}

bool append_chapter(AVFormatContext* output, std::int64_t start_us, std::int64_t end_us,
                    const std::string& title, std::int64_t id) {
  auto** chapters = static_cast<AVChapter**>(av_realloc_array(
      output->chapters, static_cast<std::size_t>(output->nb_chapters) + 1, sizeof(*output->chapters)));
  if (!chapters) return false;
  output->chapters = chapters;
  auto* chapter = static_cast<AVChapter*>(av_mallocz(sizeof(AVChapter)));
  if (!chapter) return false;
  chapter->id = id;
  chapter->time_base = AVRational{1, 1000};
  chapter->start = av_rescale_q(start_us, AV_TIME_BASE_Q, chapter->time_base);
  chapter->end = av_rescale_q(std::max(end_us, start_us + 1), AV_TIME_BASE_Q, chapter->time_base);
  if (!title.empty()) av_dict_set(&chapter->metadata, "title", title.c_str(), 0);
  output->chapters[output->nb_chapters++] = chapter;
  return true;
}
#endif
}  // namespace

std::string join_media_files(const std::vector<std::filesystem::path>& inputs,
                             const std::filesystem::path& output_path) {
#ifndef SMARTCUT_HAS_FFMPEG
  (void)inputs; (void)output_path;
  return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"FFmpeg unavailable\"}";
#else
  if (inputs.empty())
    return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"joiner list is empty\"}";
  const auto container = output_container_for_path(output_path.string());
  if (container == OutputContainer::unsupported || container == OutputContainer::elementary)
    return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"unsupported Joiner output container\"}";
  if (container != OutputContainer::matroska) {
    const auto nonce = std::chrono::high_resolution_clock::now().time_since_epoch().count();
    const auto stage = output_path.parent_path() /
                       (output_path.stem().string() + ".smartcut-join-stage-" +
                        std::to_string(nonce) + ".mkv");
    const auto join_report = join_media_files(inputs, stage);
    const bool joined = join_report.find("\"status\":\"PASSED\"") != std::string::npos;
    const auto remux_report = joined ? remux_media_container(stage.string(), output_path.string())
                                     : std::string{"{\"schema\":\"container-remux-v1\",\"status\":\"FAILED\",\"reason\":\"Joiner stage failed\"}"};
    std::error_code ignored;
    std::filesystem::remove(stage, ignored);
    const bool passed = joined && remux_report.find("\"status\":\"PASSED\"") != std::string::npos;
    if (!passed) std::filesystem::remove(output_path, ignored);
    std::ostringstream report;
    report << "{\"schema\":\"media-joiner-container-v1\",\"status\":\""
           << (passed ? "PASSED" : "FAILED") << "\",\"container\":\""
           << output_container_name(container) << "\",\"smart_copy\":true,\"join_report\":"
           << join_report << ",\"remux_report\":" << remux_report << '}';
    return report.str();
  }
  AVFormatContext* output = nullptr;
  if (avformat_alloc_output_context2(&output, nullptr, "matroska", output_path.string().c_str()) < 0 || !output)
    return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"output context failed\"}";
  auto cleanup = [&] { if (output) { if (!(output->oformat->flags & AVFMT_NOFILE) && output->pb) avio_closep(&output->pb); avformat_free_context(output); output = nullptr; } };
  std::vector<InputTiming> timings;
  std::vector<std::vector<ChapterCopy>> chapters;
  timings.reserve(inputs.size());
  chapters.reserve(inputs.size());
  std::size_t template_index = 0;
  for (std::size_t input_index = 0; input_index < inputs.size(); ++input_index) {
    auto candidate = ffmpeg::open_input(inputs[input_index]);
    if (!candidate || avformat_find_stream_info(candidate.get(), nullptr) < 0) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"input open failed\"}"; }
    const auto* title = av_dict_get(candidate->metadata, "smartcut_title", nullptr, 0);
    if (!title || std::string_view(title->value) != "1") { template_index = input_index; break; }
  }
  auto stream_template = ffmpeg::open_input(inputs[template_index]);
  if (!stream_template || avformat_find_stream_info(stream_template.get(), nullptr) < 0) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"stream template open failed\"}"; }
  const unsigned expected_streams = stream_template->nb_streams;
  av_dict_copy(&output->metadata, stream_template->metadata, 0);
  av_dict_set(&output->metadata, "smartcut_title", nullptr, 0);
  for (unsigned stream_index = 0; stream_index < expected_streams; ++stream_index) {
    auto* destination = avformat_new_stream(output, nullptr);
    const auto* source = stream_template->streams[stream_index];
    if (!destination || avcodec_parameters_copy(destination->codecpar, source->codecpar) < 0) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"stream setup failed\"}"; }
    destination->codecpar->codec_tag = 0;
    destination->time_base = source->time_base;
    destination->avg_frame_rate = source->avg_frame_rate;
    destination->r_frame_rate = source->r_frame_rate;
    destination->disposition = source->disposition;
    destination->sample_aspect_ratio = source->sample_aspect_ratio;
    av_dict_copy(&destination->metadata, source->metadata, 0);
  }
  std::int64_t cumulative_us = 0;
  for (std::size_t input_index = 0; input_index < inputs.size(); ++input_index) {
    auto input = ffmpeg::open_input(inputs[input_index]);
    if (!input || avformat_find_stream_info(input.get(), nullptr) < 0) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"input open failed\"}"; }
    if (input->nb_streams != expected_streams) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"stream counts differ\"}"; }
    for (unsigned stream_index = 0; stream_index < expected_streams; ++stream_index) {
      if (!compatible_streams(output->streams[stream_index]->codecpar, input->streams[stream_index]->codecpar)) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"stream codecs or settings differ\"}"; }
    }
    std::vector<ChapterCopy> source_chapters;
    for (unsigned chapter_index = 0; chapter_index < input->nb_chapters; ++chapter_index) {
      const auto* chapter = input->chapters[chapter_index];
      const auto start = av_rescale_q(chapter->start, chapter->time_base, AV_TIME_BASE_Q);
      const auto end = av_rescale_q(chapter->end, chapter->time_base, AV_TIME_BASE_Q);
      const AVDictionaryEntry* title = av_dict_get(chapter->metadata, "title", nullptr, 0);
      source_chapters.push_back({start, end, title ? title->value : ""});
    }
    chapters.push_back(std::move(source_chapters));
    timings.push_back(scan_timing(input.get()));
  }
  const auto segment_base_us = [](const InputTiming& timing) {
    return timing.video_first_us == AV_NOPTS_VALUE ? timing.first_us : timing.video_first_us;
  };
  const auto segment_duration_us = [&](const InputTiming& timing) {
    const auto base = segment_base_us(timing);
    const auto end = timing.video_end_us == AV_NOPTS_VALUE ? timing.end_us : timing.video_end_us;
    return std::max<std::int64_t>(end - base, 1);
  };
  const auto initial_video_lead_us =
      std::max<std::int64_t>(segment_base_us(timings.front()) - timings.front().first_us, 0);
  cumulative_us = initial_video_lead_us;
  std::uint64_t joined_video_frames = 0;
  for (const auto& timing : timings) joined_video_frames += timing.video_packets;
  for (unsigned stream_index = 0; stream_index < expected_streams; ++stream_index) {
    auto* stream = output->streams[stream_index];
    if (stream->codecpar->codec_type != AVMEDIA_TYPE_VIDEO ||
        (stream->disposition & AV_DISPOSITION_ATTACHED_PIC)) continue;
    for (const char* key : {"BPS-eng", "DURATION-eng", "NUMBER_OF_FRAMES-eng",
                            "NUMBER_OF_BYTES-eng", "_STATISTICS_WRITING_APP-eng",
                            "_STATISTICS_WRITING_DATE_UTC-eng", "_STATISTICS_TAGS-eng"})
      av_dict_set(&stream->metadata, key, nullptr, 0);
    const auto frame_count = std::to_string(joined_video_frames);
    const auto rate_num = std::to_string(stream_template->streams[stream_index]->avg_frame_rate.num);
    const auto rate_den = std::to_string(stream_template->streams[stream_index]->avg_frame_rate.den);
    av_dict_set(&stream->metadata, "SMARTCUT_FRAME_COUNT", frame_count.c_str(), 0);
    av_dict_set(&stream->metadata, "SMARTCUT_FRAME_RATE_NUM", rate_num.c_str(), 0);
    av_dict_set(&stream->metadata, "SMARTCUT_FRAME_RATE_DEN", rate_den.c_str(), 0);
  }
  for (std::size_t input_index = 0; input_index < inputs.size(); ++input_index) {
    const auto base_us = segment_base_us(timings[input_index]);
    const auto duration = segment_duration_us(timings[input_index]);
    for (const auto& chapter : chapters[input_index]) {
      const auto relative_start = std::max<std::int64_t>(chapter.start_us - base_us, 0);
      const auto relative_end = std::max<std::int64_t>(chapter.end_us - base_us, relative_start + 1);
      if (!append_chapter(output, cumulative_us + relative_start, cumulative_us + relative_end,
                          chapter.title, static_cast<std::int64_t>(output->nb_chapters))) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"chapter setup failed\"}"; }
    }
    cumulative_us += duration;
  }
  if ((!(output->oformat->flags & AVFMT_NOFILE) && avio_open(&output->pb, output_path.string().c_str(), AVIO_FLAG_WRITE) < 0) ||
      avformat_write_header(output, nullptr) < 0) { cleanup(); return "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"output header failed\"}"; }
  cumulative_us = initial_video_lead_us;
  std::uint64_t packets_written = 0;
  bool write_failed = false;
  int write_error = 0;
  std::size_t failed_input = 0;
  int failed_stream = -1;
  std::int64_t failed_pts = AV_NOPTS_VALUE;
  std::int64_t failed_dts = AV_NOPTS_VALUE;
  std::int64_t failed_base_us = 0;
  std::int64_t failed_cumulative_us = 0;
  std::int64_t failed_video_first_us = AV_NOPTS_VALUE;
  std::int64_t failed_video_end_us = AV_NOPTS_VALUE;
  int failed_media_type = AVMEDIA_TYPE_UNKNOWN;
  int failed_disposition = 0;
  std::vector<std::int64_t> last_output_dts(expected_streams, AV_NOPTS_VALUE);
  for (std::size_t input_index = 0; input_index < inputs.size() && !write_failed; ++input_index) {
    auto input = ffmpeg::open_input(inputs[input_index]);
    if (!input || avformat_find_stream_info(input.get(), nullptr) < 0) { write_failed = true; break; }
    auto packet = ffmpeg::make_packet();
    std::vector<std::int64_t> decode_shifts(expected_streams, AV_NOPTS_VALUE);
    while (packet && av_read_frame(input.get(), packet.get()) >= 0) {
      const auto stream_index = static_cast<unsigned>(packet->stream_index);
      const auto input_time_base = input->streams[stream_index]->time_base;
      const auto output_time_base = output->streams[stream_index]->time_base;
      const auto input_base_us = segment_base_us(timings[input_index]);
      const bool missing_dts = packet->dts == AV_NOPTS_VALUE;
      const auto translate = [&](std::int64_t timestamp) {
        if (timestamp == AV_NOPTS_VALUE) return timestamp;
        const auto source_us = av_rescale_q(timestamp, input_time_base, AV_TIME_BASE_Q);
        return av_rescale_q(source_us - input_base_us + cumulative_us,
                            AV_TIME_BASE_Q, output_time_base);
      };
      packet->pts = translate(packet->pts);
      packet->dts = translate(packet->dts);
      packet->duration = av_rescale_q(packet->duration, input_time_base, output_time_base);
      const bool matroska_video_packet =
          input->streams[stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO &&
          !(input->streams[stream_index]->disposition & AV_DISPOSITION_ATTACHED_PIC);
      if (matroska_video_packet) packet->dts = AV_NOPTS_VALUE;
      if (!matroska_video_packet && input_index > 0 && decode_shifts[stream_index] == AV_NOPTS_VALUE) {
        const auto step = std::max<std::int64_t>(packet->duration, 1);
        std::int64_t inferred_dts = packet->dts;
        if (inferred_dts == AV_NOPTS_VALUE) {
          const auto first_valid_us = timings[input_index].first_dts_us[stream_index];
          if (first_valid_us != AV_NOPTS_VALUE) {
            const auto inferred_us = first_valid_us - av_rescale_q(step, output_time_base, AV_TIME_BASE_Q);
            inferred_dts = av_rescale_q(inferred_us - input_base_us + cumulative_us,
                                        AV_TIME_BASE_Q, output_time_base);
          } else {
            inferred_dts = packet->pts == AV_NOPTS_VALUE ? 0 : packet->pts;
          }
        }
        const auto target_dts = last_output_dts[stream_index] == AV_NOPTS_VALUE
                                    ? inferred_dts
                                    : last_output_dts[stream_index] + step;
        decode_shifts[stream_index] = target_dts - inferred_dts;
      }
      if (!matroska_video_packet && input_index > 0) {
        if (missing_dts) {
          const auto step = std::max<std::int64_t>(packet->duration, 1);
          const auto first_valid_us = timings[input_index].first_dts_us[stream_index];
          const auto inferred_us = first_valid_us == AV_NOPTS_VALUE
                                       ? input_base_us
                                       : first_valid_us - av_rescale_q(step, output_time_base, AV_TIME_BASE_Q);
          const auto inferred_dts = av_rescale_q(inferred_us - input_base_us + cumulative_us,
                                                 AV_TIME_BASE_Q, output_time_base);
          packet->dts = inferred_dts + decode_shifts[stream_index];
        } else {
          packet->dts += decode_shifts[stream_index];
        }
        if (packet->pts != AV_NOPTS_VALUE) packet->pts += decode_shifts[stream_index];
      }
      packet->pos = -1;
      const auto submitted_pts = packet->pts;
      const auto submitted_dts = packet->dts;
      if (const auto packet_result = av_interleaved_write_frame(output, packet.get()); packet_result < 0) {
        write_failed = true;
        write_error = packet_result;
        failed_input = input_index;
        failed_stream = packet->stream_index;
        failed_pts = submitted_pts;
        failed_dts = submitted_dts;
        failed_base_us = input_base_us;
        failed_cumulative_us = cumulative_us;
        failed_video_first_us = timings[input_index].video_first_us;
        failed_video_end_us = timings[input_index].video_end_us;
        failed_media_type = input->streams[stream_index]->codecpar->codec_type;
        failed_disposition = input->streams[stream_index]->disposition;
        break;
      }
      if (submitted_dts != AV_NOPTS_VALUE) last_output_dts[stream_index] = submitted_dts;
      ++packets_written;
      av_packet_unref(packet.get());
    }
    cumulative_us += segment_duration_us(timings[input_index]);
  }
  const bool trailer_written = !write_failed && av_write_trailer(output) >= 0;
  cleanup();
  if (!trailer_written) {
    std::error_code error;
    std::filesystem::remove(output_path, error);
    char error_text[AV_ERROR_MAX_STRING_SIZE]{};
    if (write_error < 0) av_strerror(write_error, error_text, sizeof(error_text));
    std::ostringstream report;
    report << "{\"schema\":\"media-joiner-v1\",\"status\":\"FAILED\",\"error\":\"packet write failed\",\"input_index\":"
           << failed_input << ",\"stream_index\":" << failed_stream << ",\"pts\":" << failed_pts
           << ",\"dts\":" << failed_dts << ",\"base_us\":" << failed_base_us
           << ",\"cumulative_us\":" << failed_cumulative_us << ",\"video_first_us\":" << failed_video_first_us
           << ",\"video_end_us\":" << failed_video_end_us << ",\"media_type\":" << failed_media_type
           << ",\"disposition\":" << failed_disposition << ",\"ffmpeg_error\":\"" << error_text << "\"}";
    return report.str();
  }
  auto validation = ffmpeg::open_input(output_path);
  const bool valid = validation && avformat_find_stream_info(validation.get(), nullptr) >= 0 &&
                     validation->nb_streams == expected_streams;
  if (!valid) { std::error_code error; std::filesystem::remove(output_path, error); }
  std::ostringstream report;
  report << "{\"schema\":\"media-joiner-v1\",\"status\":\"" << (valid ? "PASSED" : "FAILED")
         << "\",\"inputs\":" << inputs.size() << ",\"streams\":" << expected_streams
         << ",\"packets_written\":" << packets_written << ",\"chapters\":"
         << (valid ? validation->nb_chapters : 0) << '}';
  return report.str();
#endif
}

}  // namespace smartcut::media
