/*
 * HEIF image base codec.
 * Copyright (c) 2024 Dirk Farin <dirk.farin@gmail.com>
 *
 * This file is part of libheif.
 *
 * libheif is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * libheif is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef LIBHEIF_TRACK_H
#define LIBHEIF_TRACK_H

#include "error.h"
#include "api_structs.h"
#include "security_limits.h"
#include "libheif/heif_plugin.h"
#include "libheif/heif_sequences.h"
#include <string>
#include <memory>
#include <utility>
#include <vector>

class HeifContext;

class HeifPixelImage;

class Chunk;

class Box_trak;


class SampleAuxInfoHelper
{
public:
  SampleAuxInfoHelper(bool interleaved = false);

  void set_aux_info_type(uint32_t aux_info_type, uint32_t aux_info_type_parameter = 0);

  Error add_sample_info(const std::vector<uint8_t>& data);

  void add_nonpresent_sample();

  void write_interleaved(const std::shared_ptr<HeifFile>& file);

  void write_all(const std::shared_ptr<Box>& parent, const std::shared_ptr<HeifFile>& file);

private:
  std::shared_ptr<Box_saiz> m_saiz;
  std::shared_ptr<Box_saio> m_saio;

  std::vector<uint8_t> m_data;

  bool m_interleaved;
};


class SampleAuxInfoReader
{
public:
  SampleAuxInfoReader(const std::shared_ptr<Box_saiz>&,
                      const std::shared_ptr<Box_saio>&,
                      const std::vector<std::shared_ptr<Chunk>>& chunks);

  heif_sample_aux_info_type get_type() const;

  Result<std::vector<uint8_t>> get_sample_info(const HeifFile* file, uint32_t sample_idx);

private:
  std::shared_ptr<Box_saiz> m_saiz;
  std::shared_ptr<Box_saio> m_saio;

  // If there is only one chunk and the SAI data sizes are constant, we do not need an offset table.
  // We just store the base offset and can directly calculate the sample offset from that.
  bool m_contiguous_and_constant_size=false;
  uint64_t m_singleChunk_offset=0;

  // For chunked data or non-constant sample sizes, we use a table with the offsets for all SAI samples.
  std::vector<uint64_t> m_sample_offsets;
};


/**
 * This structure specifies what will be written in a track and how it will be laid out in the file.
 */
struct TrackOptions
{
  ~TrackOptions()
  {
    heif_tai_clock_info_release(tai_clock_info);
  }

  // Timescale (clock ticks per second) for this track.
  uint32_t track_timescale = 90000;

  // If 'true', the aux_info data blocks will be interleaved with the compressed image.
  // This has the advantage that the aux_info is localized near the image data.
  //
  // If 'false', all aux_info will be written as one block after the compressed image data.
  // This has the advantage that no aux_info offsets have to be written.
  bool write_sample_aux_infos_interleaved = false;


  // --- TAI timestamps for samples
  heif_sample_aux_info_presence with_sample_tai_timestamps = heif_sample_aux_info_presence_none;
  heif_tai_clock_info* tai_clock_info = nullptr;

  // --- GIMI content IDs for samples

  heif_sample_aux_info_presence with_sample_content_ids = heif_sample_aux_info_presence_none;

  // --- GIMI content ID for the track

  std::string gimi_track_content_id;

  TrackOptions& operator=(const TrackOptions&);
};


const char* get_track_auxiliary_info_type(heif_compression_format format);


class Track : public ErrorBuffer {
public:
  //Track(HeifContext* ctx);

  Track(HeifContext* ctx, uint32_t track_id, const TrackOptions* info, uint32_t handler_type);

  Track(HeifContext* ctx);

  virtual ~Track() = default;

  // Allocate a Track of the correct sub-class (visual or metadata).
  // For tracks with an unsupported handler type, heif_error_Unsupported_feature/heif_suberror_Unsupported_track_type is returned.
  static Result<std::shared_ptr<Track>> alloc_track(HeifContext*, const std::shared_ptr<Box_trak>&);

  // load track from file
  virtual Error load(const std::shared_ptr<Box_trak>&);

  // This is called after creating all Track objects when reading a HEIF file.
  // We can now do initializations that require access to all tracks.
  [[nodiscard]] virtual Error initialize_after_parsing(HeifContext*, const std::vector<std::shared_ptr<Track>>& all_tracks) { return {}; }

  heif_item_id get_id() const { return m_id; }

  std::shared_ptr<HeifFile> get_file() const;

  uint32_t get_handler() const { return m_handler_type; }

  heif_auxiliary_track_info_type get_auxiliary_info_type() const;

  std::string get_auxiliary_info_type_urn() const { return m_auxiliary_info_type; }

  void set_auxiliary_info_type(heif_auxiliary_track_info_type);

  void set_auxiliary_info_type_urn(std::string t) { m_auxiliary_info_type = std::move(t); }

  bool is_visual_track() const;

  virtual bool has_alpha_channel() const { return false; }

  uint32_t get_first_cluster_sample_entry_type() const;

  Result<std::string> get_first_cluster_urim_uri() const;

  uint64_t get_duration_in_media_units() const;

  uint32_t get_timescale() const;

  // The context will compute the duration in global movie units and set this.
  void set_track_duration_in_movie_units(uint64_t total_duration, uint64_t segment_duration);

  void enable_edit_list_repeat_mode(bool enable);

  std::shared_ptr<Box_taic> get_first_cluster_taic() { return m_first_taic; }

  bool end_of_sequence_reached() const;

  // See m_num_repetitions for the meaning of the return value.
  uint32_t get_number_of_repetitions() const { return m_num_repetitions; }

  // Compute some parameters after all frames have been encoded (for example: track duration).
  virtual Error finalize_track();

  const TrackOptions& get_track_info() const { return m_track_info; }

  void add_reference_to_track(uint32_t referenceType, uint32_t to_track_id);

  std::shared_ptr<const Box_tref> get_tref_box() const { return m_tref; }

  Result<heif_raw_sequence_sample*> get_next_sample_raw_data(const heif_decoding_options* options);

  std::vector<heif_sample_aux_info_type> get_sample_aux_info_types() const;

protected:
  HeifContext* m_heif_context = nullptr;
  uint32_t m_id = 0;
  uint32_t m_handler_type = 0;

  TrackOptions m_track_info;

  uint32_t m_num_samples = 0;

  struct SampleTiming {
    uint32_t sampleIdx = 0;
    uint32_t sampleInChunkIdx = 0;
    uint32_t chunkIdx = 0;
    uint64_t presentation_time = 0; // TODO
    uint64_t media_composition_time = 0; // TODO
    uint64_t media_decoding_time = 0;
    uint32_t sample_duration_media_time = 0;
    uint32_t sample_duration_presentation_time = 0; // TODO
  };
  std::vector<SampleTiming> m_presentation_timeline;
  // Accounts m_presentation_timeline against max_total_memory. A small track can
  // declare millions of samples, so this vector (~48 bytes/sample) must be tracked
  // to bound multi-track accumulation (GHSA-xw34-mjcp-jqh8, variants V2/V3).
  MemoryHandle m_presentation_timeline_memory;
  uint64_t m_num_output_samples = 0; // Can be larger than the vector. It then repeats the playback.

  // How many times the media timeline is repeated.
  // 0  = editlist is present but its pattern is not understood (caller should assume a single playback).
  // 1  = no editlist: media plays exactly once.
  // UINT32_MAX = infinite (mvhd duration is the indefinite-sentinel and the editlist is in repeat mode).
  // N  = the media segment is played N times.
  uint32_t m_num_repetitions = 1;

  // Continuous counting through all repetitions. You have to take the modulo operation to get the
  // index into m_presentation_timeline SampleTiming table.
  // (At 30 fps, this 32 bit integer will overflow in >4 years. I think this is acceptable.)
  uint32_t m_next_sample_to_be_decoded = 0;

  // Total sequence output index.
  uint32_t m_next_sample_to_be_output = 0;
  bool     m_decoder_is_flushed = false;

  Error init_sample_timing_table();

  std::vector<std::shared_ptr<Chunk>> m_chunks;
  // Accounts the per-chunk Chunk::m_sample_ranges tables against max_total_memory.
  // Their combined size over all chunks is num_samples * sizeof(SampleFileRange),
  // so a single reservation here bounds them all (GHSA-xw34-mjcp-jqh8, V2/V3).
  MemoryHandle m_chunk_sample_ranges_memory;
  std::vector<uint8_t> m_chunk_data;

  std::shared_ptr<Box_moov> m_moov;
  std::shared_ptr<Box_trak> m_trak;
  std::shared_ptr<Box_tkhd> m_tkhd;
  std::shared_ptr<Box_minf> m_minf;
  std::shared_ptr<Box_mdhd> m_mdhd;
  std::shared_ptr<Box_hdlr> m_hdlr;
  std::shared_ptr<Box_stbl> m_stbl;
  std::shared_ptr<Box_stsd> m_stsd;
  std::shared_ptr<Box_stsc> m_stsc;
  std::shared_ptr<Box_stco> m_stco;
  std::shared_ptr<Box_stts> m_stts;
  std::shared_ptr<Box_ctts> m_ctts; // optional box, TODO: add only if needed
  std::shared_ptr<Box_stss> m_stss;
  std::shared_ptr<Box_stsz> m_stsz;
  std::shared_ptr<Box_elst> m_elst;

  std::shared_ptr<class Box_tref> m_tref; // optional

  std::string m_auxiliary_info_type; // only for auxiliary tracks

  // --- sample auxiliary information

  std::unique_ptr<SampleAuxInfoHelper> m_aux_helper_tai_timestamps;
  std::unique_ptr<SampleAuxInfoHelper> m_aux_helper_content_ids;

  std::unique_ptr<SampleAuxInfoReader> m_aux_reader_tai_timestamps;
  std::unique_ptr<SampleAuxInfoReader> m_aux_reader_content_ids;

  std::shared_ptr<class Box_taic> m_first_taic; // the TAIC of the first chunk


  // --- Helper functions for writing samples.

  // Call when we begin a new chunk of samples, e.g. because the compression format changed
  void add_chunk(heif_compression_format format);

  // Call to set the sample_description_box for the last added chunk.
  // Has to be called when we call add_chunk().
  // It is not merged with add_chunk() because the sample_description_box may need information from the
  // first encoded frame.
  void set_sample_description_box(const std::shared_ptr<Box>& sample_description_box);

  // Write the actual sample data. `tai` may be null and `gimi_contentID` may be empty.
  // In these cases, no timestamp or no contentID will be written, respectively.
  Error write_sample_data(const std::vector<uint8_t>& raw_data,
                          uint32_t sample_duration,
                          int32_t composition_time_offset,
                          bool is_sync_sample,
                          const heif_tai_timestamp_packet* tai,
                          const std::optional<std::string>& gimi_contentID);
};


#endif //LIBHEIF_TRACK_H
