/*
   mkvextract -- extract tracks from Matroska files into other files

   Distributed under the GPL
   see the file COPYING for details
   or visit http://www.gnu.org/copyleft/gpl.html

   extracts tracks from Matroska files into other files

   Written by Moritz Bunkus <moritz@bunkus.org>.
*/

#include "common/common_pch.h"

#include <cassert>

#include <ebml/EbmlHead.h>
#include <ebml/EbmlSubHead.h>
#include <ebml/EbmlStream.h>
#include <ebml/EbmlVersion.h>
#include <ebml/EbmlVoid.h>

#include <matroska/KaxBlock.h>
#include <matroska/KaxBlockData.h>
#include <matroska/KaxCluster.h>
#include <matroska/KaxClusterData.h>
#include <matroska/KaxInfo.h>
#include <matroska/KaxInfoData.h>
#include <matroska/KaxSegment.h>
#include <matroska/KaxTracks.h>
#include <matroska/KaxTrackEntryData.h>
#include <matroska/KaxTrackAudio.h>
#include <matroska/KaxTrackVideo.h>

#include "common/ebml.h"
#include "common/kax_file.h"
#include "common/mm_io_x.h"
#include "common/mm_write_buffer_io.h"
#include "extract/mkvextract.h"
#include "extract/xtr_base.h"

using namespace libmatroska;

static std::vector<xtr_base_c *> extractors;

// ------------------------------------------------------------------------

static void
create_extractors(KaxTracks &kax_tracks,
                  std::vector<track_spec_t> &tracks) {
  size_t i;
  int64_t track_id = -1;

  for (i = 0; i < kax_tracks.ListSize(); i++) {
    if (!Is<KaxTrackEntry>(kax_tracks[i]))
      continue;

    ++track_id;

    KaxTrackEntry &track = *static_cast<KaxTrackEntry *>(kax_tracks[i]);
    int64_t tnum         = kt_get_number(track);

    // Is the track number present and valid?
    if (0 == tnum)
      continue;

    // Is there more than one track with the same track number?
    xtr_base_c *extractor = nullptr;
    size_t k;
    for (k = 0; k < extractors.size(); k++)
      if (extractors[k]->m_track_num == tnum) {
        mxwarn(boost::format(Y("More than one track with the track number %1% found.\n")) % tnum);
        extractor = extractors[k];
        break;
      }
    if (extractor)
      continue;

    // Does the user want this track to be extracted?
    track_spec_t *tspec = nullptr;
    for (k = 0; k < tracks.size(); k++)
      if (tracks[k].tid == track_id) {
        tspec = &tracks[k];
        break;
      }
    if (!tspec)
      continue;

    // Let's find the codec ID and create an extractor for it.
    std::string codec_id = kt_get_codec_id(track);
    if (codec_id.empty())
      mxerror(boost::format(Y("The track ID %1% does not have a valid CodecID.\n")) % track_id);

    extractor = xtr_base_c::create_extractor(codec_id, track_id, *tspec);
    if (!extractor)
      mxerror(boost::format(Y("Extraction of track ID %1% with the CodecID '%2%' is not supported.\n")) % track_id % codec_id);

    extractor->m_track_num = tnum;

    // Has there another file been requested with the same name?
    xtr_base_c *master = nullptr;
    for (k = 0; k < extractors.size(); k++)
      if (extractors[k]->m_file_name == tspec->out_name) {
        master = extractors[k];
        break;
      }

    // Let the extractor create the file.
    extractor->create_file(master, track);

    // We're done.
    extractors.push_back(extractor);

    mxinfo(boost::format(Y("Extracting track %1% with the CodecID '%2%' to the file '%3%'. Container format: %4%\n"))
           % track_id % codec_id % tspec->out_name % extractor->get_container_name());
  }

  // Signal that all headers have been taken care of.
  for (i = 0; i < extractors.size(); i++)
    extractors[i]->headers_done();
}

static int64_t
handle_blockgroup(KaxBlockGroup &blockgroup,
                  KaxCluster &cluster,
                  int64_t tc_scale) {
  // Only continue if this block group actually contains a block.
  KaxBlock *block = FindChild<KaxBlock>(&blockgroup);
  if (!block || (0 == block->NumberFrames()))
    return -1;

  block->SetParent(cluster);

  // Do we need this block group?
  xtr_base_c *extractor = nullptr;
  size_t i;
  for (i = 0; i < extractors.size(); i++)
    if (block->TrackNum() == extractors[i]->m_track_num) {
      extractor = extractors[i];
      break;
    }
  if (!extractor)
    return -1;

  // Next find the block duration if there is one.
  KaxBlockDuration *kduration   = FindChild<KaxBlockDuration>(&blockgroup);
  int64_t duration              = !kduration ? -1 : static_cast<int64_t>(kduration->GetValue() * tc_scale);
  int64_t max_timecode          = 0;

  // Now find backward and forward references.
  int64_t bref    = 0;
  int64_t fref    = 0;
  auto kreference = FindChild<KaxReferenceBlock>(&blockgroup);
  for (i = 0; (2 > i) && kreference; i++) {
    if (0 > kreference->GetValue())
      bref = kreference->GetValue();
    else
      fref = kreference->GetValue();
    kreference = FindNextChild<KaxReferenceBlock>(&blockgroup, kreference);
  }

  // Any block additions present?
  KaxBlockAdditions *kadditions = FindChild<KaxBlockAdditions>(&blockgroup);

  if (0 > duration)
    duration = extractor->m_default_duration * block->NumberFrames();

  KaxCodecState *kcstate = FindChild<KaxCodecState>(&blockgroup);
  if (kcstate) {
    memory_cptr codec_state(new memory_c(kcstate->GetBuffer(), kcstate->GetSize(), false));
    extractor->handle_codec_state(codec_state);
  }

  for (i = 0; i < block->NumberFrames(); i++) {
    int64_t this_timecode, this_duration;

    if (0 > duration) {
      this_timecode = block->GlobalTimecode();
      this_duration = duration;
    } else {
      this_timecode = block->GlobalTimecode() + i * duration /
        block->NumberFrames();
      this_duration = duration / block->NumberFrames();
    }

    auto discard_padding  = timecode_c::ns(0);
    auto kdiscard_padding = FindChild<KaxDiscardPadding>(blockgroup);
    if (kdiscard_padding)
      discard_padding = timecode_c::ns(kdiscard_padding->GetValue());

    auto &data = block->GetBuffer(i);
    auto frame = std::make_shared<memory_c>(data.Buffer(), data.Size(), false);
    auto f     = xtr_frame_t{frame, kadditions, this_timecode, this_duration, bref, fref, false, false, true, discard_padding};
    extractor->decode_and_handle_frame(f);

    max_timecode = std::max(max_timecode, this_timecode);
  }

  return max_timecode;
}

static int64_t
handle_simpleblock(KaxSimpleBlock &simpleblock,
                   KaxCluster &cluster) {
  if (0 == simpleblock.NumberFrames())
    return - 1;

  simpleblock.SetParent(cluster);

  // Do we need this block group?
  xtr_base_c *extractor = nullptr;
  size_t i;
  for (i = 0; i < extractors.size(); i++)
    if (simpleblock.TrackNum() == extractors[i]->m_track_num) {
      extractor = extractors[i];
      break;
    }

  if (!extractor)
    return - 1;

  int64_t duration     = extractor->m_default_duration * simpleblock.NumberFrames();
  int64_t max_timecode = 0;

  for (i = 0; i < simpleblock.NumberFrames(); i++) {
    int64_t this_timecode, this_duration;

    if (0 > duration) {
      this_timecode = simpleblock.GlobalTimecode();
      this_duration = duration;
    } else {
      this_timecode = simpleblock.GlobalTimecode() + i * duration / simpleblock.NumberFrames();
      this_duration = duration / simpleblock.NumberFrames();
    }

    auto &data = simpleblock.GetBuffer(i);
    auto frame = std::make_shared<memory_c>(data.Buffer(), data.Size(), false);
    auto f     = xtr_frame_t{frame, nullptr, this_timecode, this_duration, -1, -1, simpleblock.IsKeyframe(), simpleblock.IsDiscardable(), false, timecode_c::ns(0)};
    extractor->decode_and_handle_frame(f);

    max_timecode = std::max(max_timecode, this_timecode);
  }

  return max_timecode;
}

static void
close_extractors() {
  size_t i;

  for (i = 0; i < extractors.size(); i++)
    extractors[i]->finish_track();

  for (i = 0; i < extractors.size(); i++)
    if (extractors[i]->m_master)
      extractors[i]->finish_file();

  for (i = 0; i < extractors.size(); i++) {
    if (!extractors[i]->m_master)
      extractors[i]->finish_file();
    delete extractors[i];
  }

  extractors.clear();
}

static void
write_all_cuesheets(KaxChapters &chapters,
                    KaxTags &tags,
                    std::vector<track_spec_t> &tspecs) {
  size_t i;

  for (i = 0; i < tspecs.size(); i++) {
    if (tspecs[i].extract_cuesheet) {
      std::string file_name = tspecs[i].out_name;
      int pos               = file_name.rfind('/');
      int pos2              = file_name.rfind('\\');

      if (pos2 > pos)
        pos = pos2;
      if (pos >= 0)
        file_name.erase(0, pos2);

      std::string cue_file_name = tspecs[i].out_name;
      pos                       = cue_file_name.rfind('.');
      pos2                      = cue_file_name.rfind('/');
      int pos3                  = cue_file_name.rfind('\\');

      if ((0 <= pos) && (pos > pos2) && (pos > pos3))
        cue_file_name.erase(pos);
      cue_file_name += ".cue";

      try {
        mm_io_cptr out = mm_write_buffer_io_c::open(cue_file_name, 128 * 1024);
        mxinfo(boost::format(Y("The CUE sheet for track %1% will be written to '%2%'.\n")) % tspecs[i].tid % cue_file_name);
        write_cuesheet(file_name, chapters, tags, tspecs[i].tuid, *out);

      } catch(mtx::mm_io::open_x &ex) {
        mxerror(boost::format(Y("The file '%1%' could not be opened for writing: %2%.\n")) % cue_file_name % ex);
      }
    }
  }
}

void
find_and_verify_track_uids(KaxTracks &tracks,
                           std::vector<track_spec_t> &tspecs) {
  std::map<int64_t, bool> available_track_ids;
  size_t t;
  int64_t track_id = -1;

  for (t = 0; t < tracks.ListSize(); t++) {
    KaxTrackEntry *track_entry = dynamic_cast<KaxTrackEntry *>(tracks[t]);
    if (!track_entry)
      continue;

    ++track_id;
    available_track_ids[track_id] = true;

    for (auto &tspec : tspecs)
      if (tspec.tid == track_id) {
        tspec.tuid = kt_get_uid(*track_entry);
        break;
      }
  }

  for (auto &tspec : tspecs)
    if (!available_track_ids[ tspec.tid ])
      mxerror(boost::format(Y("No track with the ID %1% was found in the source file.\n")) % tspec.tid);
}

static void
handle_segment_info(EbmlMaster *info,
                    kax_file_c *file,
                    uint64_t &tc_scale) {
  auto ktc_scale = FindChild<KaxTimecodeScale>(info);
  if (!ktc_scale)
    return;

  tc_scale = ktc_scale->GetValue();
  file->set_timecode_scale(tc_scale);
}

bool
extract_tracks(const std::string &file_name,
               std::vector<track_spec_t> &tspecs,
               kax_analyzer_c::parse_mode_e parse_mode) {
  if (tspecs.empty())
    mxerror(Y("Nothing to do.\n"));

  // open input file
  mm_io_cptr in;
  kax_file_cptr file;
  try {
    in   = mm_file_io_c::open(file_name);
    file = kax_file_cptr(new kax_file_c(in));
  } catch (mtx::mm_io::exception &ex) {
    show_error(boost::format(Y("The file '%1%' could not be opened for reading: %2%.\n")) % file_name % ex);
    return false;
  }

  int64_t file_size = in->get_size();
  uint64_t tc_scale = TIMECODE_SCALE;
  bool segment_info_found = false, tracks_found = false;

  // open input file
  auto analyzer = std::make_shared<kax_analyzer_c>(static_cast<mm_file_io_c *>(in.get()));
  if (analyzer->process(parse_mode, MODE_READ)) {
    auto af_master    = ebml_master_cptr{ analyzer->read_all(EBML_INFO(KaxInfo)) };
    auto segment_info = dynamic_cast<KaxInfo *>(af_master.get());
    if (segment_info) {
      segment_info_found = true;
      handle_segment_info(segment_info, file.get(), tc_scale);
    }

    af_master   = ebml_master_cptr{ analyzer->read_all(EBML_INFO(KaxTracks)) };
    auto tracks = dynamic_cast<KaxTracks *>(af_master.get());
    if (tracks) {
      tracks_found = true;
      find_and_verify_track_uids(*tracks, tspecs);
      create_extractors(*tracks, tspecs);
    }
  }

  try {
    in->setFilePointer(0);
    EbmlStream *es = new EbmlStream(*in);

    // Find the EbmlHead element. Must be the first one.
    EbmlElement *l0 = es->FindNextID(EBML_INFO(EbmlHead), 0xFFFFFFFFL);
    if (!l0) {
      show_error(Y("Error: No EBML head found."));
      delete es;

      return false;
    }

    // Don't verify its data for now.
    l0->SkipData(*es, EBML_CONTEXT(l0));
    delete l0;

    while (1) {
      // Next element must be a segment
      l0 = es->FindNextID(EBML_INFO(KaxSegment), 0xFFFFFFFFFFFFFFFFLL);

      if (!l0) {
        show_error(Y("No segment/level 0 element found."));
        return false;
      }

      if (Is<KaxSegment>(l0)) {
        show_element(l0, 0, Y("Segment"));
        break;
      }

      l0->SkipData(*es, EBML_CONTEXT(l0));
      delete l0;
    }

    EbmlElement *l1   = nullptr;

    KaxChapters all_chapters;
    KaxTags all_tags;

    while ((l1 = file->read_next_level1_element())) {
      if (Is<KaxInfo>(l1) && !segment_info_found) {
        segment_info_found = true;
        handle_segment_info(static_cast<EbmlMaster *>(l1), file.get(), tc_scale);

      } else if (Is<KaxTracks>(l1) && !tracks_found) {
        tracks_found = true;
        find_and_verify_track_uids(*dynamic_cast<KaxTracks *>(l1), tspecs);
        create_extractors(*dynamic_cast<KaxTracks *>(l1), tspecs);

      } else if (Is<KaxCluster>(l1)) {
        show_element(l1, 1, Y("Cluster"));
        KaxCluster *cluster = static_cast<KaxCluster *>(l1);

        if (0 == verbose)
          mxinfo(boost::format(Y("Progress: %1%%%%2%")) % (int)(in->getFilePointer() * 100 / file_size) % "\r");

        KaxClusterTimecode *ctc = FindChild<KaxClusterTimecode>(l1);
        if (ctc) {
          uint64_t cluster_tc = ctc->GetValue();
          show_element(ctc, 2, boost::format(Y("Cluster timecode: %|1$.3f|s")) % ((float)cluster_tc * (float)tc_scale / 1000000000.0));
          cluster->InitTimecode(cluster_tc, tc_scale);
        } else
          cluster->InitTimecode(0, tc_scale);

        size_t i;
        int64_t max_timecode = -1;

        for (i = 0; cluster->ListSize() > i; ++i) {
          int64_t max_bg_timecode = -1;
          EbmlElement *el         = (*cluster)[i];

          if (Is<KaxBlockGroup>(el)) {
            show_element(el, 2, Y("Block group"));
            max_bg_timecode = handle_blockgroup(*static_cast<KaxBlockGroup *>(el), *cluster, tc_scale);

          } else if (Is<KaxSimpleBlock>(el)) {
            show_element(el, 2, Y("SimpleBlock"));
            max_bg_timecode = handle_simpleblock(*static_cast<KaxSimpleBlock *>(el), *cluster);
          }

          max_timecode = std::max(max_timecode, max_bg_timecode);
        }

        if (-1 != max_timecode)
          file->set_last_timecode(max_timecode);

      } else if (Is<KaxChapters>(l1)) {
        KaxChapters &chapters = *static_cast<KaxChapters *>(l1);

        while (chapters.ListSize() > 0) {
          if (Is<KaxEditionEntry>(chapters[0])) {
            KaxEditionEntry &entry = *static_cast<KaxEditionEntry *>(chapters[0]);
            while (entry.ListSize() > 0) {
              if (Is<KaxChapterAtom>(entry[0]))
                all_chapters.PushElement(*entry[0]);
              entry.Remove(0);
            }
          }
          chapters.Remove(0);
        }

      } else if (Is<KaxTags>(l1)) {
        KaxTags &tags = *static_cast<KaxTags *>(l1);

        while (tags.ListSize() > 0) {
          all_tags.PushElement(*tags[0]);
          tags.Remove(0);
        }

      }

      delete l1;

    } // while (l1)

    delete l0;
    delete es;

    write_all_cuesheets(all_chapters, all_tags, tspecs);

    // Now just close the files and go to sleep. Mummy will sing you a
    // lullaby. Just close your eyes, listen to her sweet voice, singing,
    // singing, fading... fad... ing...
    close_extractors();

    return true;
  } catch (...) {
    show_error(Y("Caught exception"));

    return false;
  }
}
