/*
 * Copyright (c) 2023, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 3-Clause Clear License
 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
 * License was not distributed with this source code in the LICENSE file, you
 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  If the
 * Alliance for Open Media Patent License 1.0 was not distributed with this
 * source code in the PATENTS file, you can obtain it at
 * aomedia.org/license/patent-license/.
 */

// Tool to dump frame data from an AVM stream as protobufs.
// Provides a superset of the functionality of
// avm/examples/inspect.c, and dumps frame data as a proto
// instead of JSON.

#include <execinfo.h>
#include <stdio.h>

#include <algorithm>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <memory>
#include <iostream>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

#include "config/aom_config.h"

#include "aom/aom_decoder.h"
#include "aom/aomdx.h"
#include "av1/common/av1_common_int.h"
#include "av1/decoder/decoder.h"
#include "av1/decoder/accounting.h"
#include "av1/decoder/inspection.h"
#include "common/args.h"
#include "common/tools_common.h"
#include "common/video_common.h"
#include "common/video_reader.h"
#include "tools/extract_proto/enum_mappings.h"

#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/log/check.h"
#include "absl/log/flags.h"
#include "absl/log/globals.h"
#include "absl/log/initialize.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "google/protobuf/text_format.h"
#include "avm_frame.pb.h"

using ::avm::tools::BlockSize;
using ::avm::tools::CodingUnit;
using ::avm::tools::EnumMappings;
using ::avm::tools::Frame;
using ::avm::tools::FrameParams;
using ::avm::tools::Partition;
using ::avm::tools::PixelBuffer;
using ::avm::tools::Position;
using ::avm::tools::StreamParams;
using ::avm::tools::Superblock;
using ::avm::tools::Symbol;
using ::avm::tools::SymbolInfo;

ABSL_FLAG(std::string, stream, "", "Input AV2 stream");
ABSL_FLAG(std::string, orig_yuv, "",
          "Source (pre-encode) YUV file (.yuv or .y4m)");
ABSL_FLAG(int, orig_yuv_bit_depth, -1, "Bit depth of original YUV.");
ABSL_FLAG(std::string, output_folder, "", "Output folder");
ABSL_FLAG(std::string, output_prefix, "",
          "Prefix added to output filenames, e.g. "
          "{output_folder}/{output_prefix}_frame_{frame_id}.pb. By default, "
          "uses the same name is the input stream.");
ABSL_FLAG(bool, output_as_text, false,
          "Proto will be output as text (.textproto) instead of binary (.pb)");
ABSL_FLAG(std::vector<std::string>, encoder_args, {},
          "Comma-separated list of encoder arguments.");
ABSL_FLAG(int, limit, -1,
          "Stop after N frames are decoded (default: all frames).");
ABSL_FLAG(bool, show_progress, true,
          "Print progress as each frame is decoded.");
ABSL_FLAG(bool, ignore_y4m_header, false,
          "Ignore dimensions and colorspace in Y4M header.");
ABSL_FLAG(int, preserve_stream_path_depth, 0,
          "Preserve this many levels of directory hierarchy of the stream's "
          "path. -1 will preserve the entire absolute path. 0 will keep only "
          "the filename. 1 will keep just the direct parent, etc...");

namespace {
constexpr std::string_view kY4mFrameMarker = "FRAME\n";
// Read ahead this many bytes when looking for the next y4m frame marker.
constexpr size_t kY4mReadahead = 1024;
constexpr std::string_view kBinaryProtoExt = "pb";
constexpr std::string_view kTextProtoExt = "textproto";

struct OutputConfig {
  std::filesystem::path output_folder;
  std::string output_prefix;
  bool output_as_text_proto;
  int limit;
};

struct Y4mHeader {
  size_t width;
  size_t height;
  size_t bit_depth;
};

static absl::flat_hash_map<std::string_view, int> kSupportedColorspaces = {
  { "C420", 8 },     { "C420jpeg", 8 }, { "C420mpeg2", 8 }, { "C420p10", 10 },
  { "C420p12", 12 }, { "C420p14", 14 }, { "C420p16", 16 },
};

struct ExtractProtoContext {
  insp_frame_data frame_data;
  aom_codec_ctx_t codec;
  AvxVideoReader *reader;
  const AvxVideoInfo *info;
  StreamParams *stream_params;
  // Note: Original YUV may not always be available (e.g. we have only an .ivf
  // file by itself)
  std::ifstream *orig_yuv_file;
  int orig_yuv_bit_depth;
  std::filesystem::path stream_path;
  int preserve_stream_path_depth;
  int decode_count;
  bool is_y4m_file;
  std::optional<Y4mHeader> y4m_header;
  OutputConfig output_config;
  bool show_progress;
  // Offset to convert relative display index to absolute display index.
  int display_index_offset;
  // Largest display index seen so far.
  int max_display_index;
};

BlockSize MakeBlockSize(BLOCK_SIZE raw_size) {
  BlockSize block_size;
  block_size.set_width(mi_size_wide[raw_size] * MI_SIZE);
  block_size.set_height(mi_size_high[raw_size] * MI_SIZE);
  block_size.set_enum_value(raw_size);
  return block_size;
}

Position MakePosition(int mi_row, int mi_col) {
  Position pos;
  pos.set_x(mi_col * MI_SIZE);
  pos.set_y(mi_row * MI_SIZE);
  return pos;
}

enum class PartitionType {
  kShared = 0,
  kLumaOnly = 1,
  kChromaOnly = 2,
};

void GetCodingUnit(CodingUnit *coding_unit, insp_frame_data *frame_data,
                   insp_sb_data *sb_data, int mi_row, int mi_col,
                   PartitionType part_type, int (&coeff_idx)[3]) {
  insp_mi_data *mi =
      &frame_data->mi_grid[mi_row * frame_data->mi_cols + mi_col];
  int sb_type = (part_type == PartitionType::kChromaOnly) ? mi->sb_type_chroma
                                                          : mi->sb_type;
  int width = mi_size_wide[sb_type];
  int height = mi_size_high[sb_type];
  coding_unit->mutable_size()->set_width(width * MI_SIZE);
  coding_unit->mutable_size()->set_height(height * MI_SIZE);
  coding_unit->mutable_position()->set_x(mi_col * MI_SIZE);
  coding_unit->mutable_position()->set_y(mi_row * MI_SIZE);
  coding_unit->set_skip(mi->skip);
  coding_unit->set_qindex(mi->current_qindex);
  coding_unit->set_segment_id(mi->segment_id);
  coding_unit->set_cdef_level(mi->cdef_level);
  coding_unit->set_cdef_strength(mi->cdef_strength);
  auto *pred = coding_unit->mutable_prediction_mode();
  pred->set_mode(mi->mode);
  pred->set_angle_delta(mi->angle_delta);
  pred->set_uv_mode(mi->uv_mode);
  pred->set_uv_angle_delta(mi->uv_angle_delta);
  pred->set_cfl_alpha_idx(mi->cfl_alpha_idx);
  pred->set_cfl_alpha_sign(mi->cfl_alpha_sign);
  pred->set_compound_type(mi->compound_type);
  pred->set_motion_mode(mi->motion_mode);
  pred->set_use_intrabc(mi->intrabc);
  // TODO(comc): Check correct enum for filter (InterpFilter: filter.h)
  pred->set_interpolation_filter(mi->filter[0]);
  // TODO(comc): Add palette colors
  pred->set_palette_count(mi->palette);
  pred->set_uv_palette_count(mi->uv_palette);
  for (int i = 0; i < 2; i++) {
    auto *mv = pred->add_motion_vectors();
    mv->set_dx(mi->mv[i].col);
    mv->set_dy(mi->mv[i].row);
    mv->set_ref_frame(mi->ref_frame[i]);
    mv->set_ref_frame_order_hint(mi->ref_frame_order_hint[i]);
    mv->set_ref_frame_is_tip(mi->ref_frame_is_tip[i]);
    mv->set_ref_frame_is_inter(mi->ref_frame_is_inter[i]);
  }
  pred->set_use_intrabc(mi->intrabc);
  pred->set_motion_vector_precision(mi->mv_precision);
  // TODO(comc): Handle transform partition trees
  int tx_size = mi->tx_size;
  int tx_height = tx_size_high_unit[tx_size];
  int tx_width = tx_size_wide_unit[tx_size];
  if (part_type == PartitionType::kChromaOnly) {
    tx_size = sb_type;
    tx_width = std::min(width, 32);
    tx_height = std::min(height, 32);
  }

  int tx_cols = width / tx_width;
  int tx_rows = height / tx_height;
  int plane_start = (part_type == PartitionType::kChromaOnly) ? 1 : 0;
  int plane_end = (part_type == PartitionType::kLumaOnly) ? 1 : 3;

  for (int plane = plane_start; plane < plane_end; plane++) {
    auto *tx_plane = coding_unit->add_transform_planes();
    tx_plane->set_plane(plane);

    for (int i = 0; i < tx_rows; i++) {
      for (int j = 0; j < tx_cols; j++) {
        int tx_mi_row = mi_row + i * tx_height;
        int tx_mi_col = mi_col + j * tx_width;
        insp_mi_data *tx_mi =
            &frame_data->mi_grid[tx_mi_row * frame_data->mi_cols + tx_mi_col];
        auto *tu = tx_plane->add_transform_units();
        tu->mutable_position()->set_y(tx_mi_row * MI_SIZE);
        tu->mutable_position()->set_x(tx_mi_col * MI_SIZE);
        tu->mutable_size()->set_height(tx_height * MI_SIZE);
        tu->mutable_size()->set_width(tx_width * MI_SIZE);
        tu->mutable_size()->set_enum_value(tx_size);
        bool skip = (tx_mi_row >= frame_data->mi_rows ||
                     tx_mi_col >= frame_data->mi_cols || tx_mi->skip);
        // For TX sizes > 32x32, all coeffs are zero except for top-left 32x32.
        int coeffs_width = std::min(32, tx_width * MI_SIZE);
        int coeffs_height = std::min(32, tx_height * MI_SIZE);
        // TODO(comc): Resolve skip/skip_mode ambiguity
        tu->set_skip(skip);
        if (!skip) {
          tu->set_tx_type(tx_mi->tx_type);
          int index = coeff_idx[plane];
          // TODO(comc): Fix transform coeffs for large blocks: Only upper 32x32
          // (or maybe this has changed) is signaled, everything else is 0.
          for (int ty = 0; ty < coeffs_width; ty++) {
            for (int tx = 0; tx < coeffs_height; tx++) {
              int dequant_val = sb_data->dequant_values[plane][index];
              tu->add_dequantizer_values(dequant_val);
              int qcoeff = sb_data->qcoeff[plane][index];
              tu->add_quantized_coeffs(qcoeff);
              int dqcoeff = sb_data->dqcoeff[plane][index];
              tu->add_dequantized_coeffs(dqcoeff);
              index += 1;
            }
          }
        }
        coeff_idx[plane] += (tx_width * MI_SIZE) * (tx_height * MI_SIZE);
      }
    }
  }
}

int PopulatePartitionTree(insp_frame_data *frame_data, insp_sb_data *sb_data,
                          Partition *partition, Superblock *sb,
                          PARTITION_TREE *tree, PartitionType part_type,
                          int (&coeff_idx)[3], int coding_unit_start) {
  partition->set_partition_type(tree->partition);
  *partition->mutable_size() = MakeBlockSize(tree->bsize);
  *partition->mutable_position() = MakePosition(tree->mi_row, tree->mi_col);
  partition->mutable_coding_unit_range()->set_start(coding_unit_start);
  // coding_unit_start = index of next coding unit to be inserted.
  int coding_unit_end = coding_unit_start;

  int child_coding_unit_start = coding_unit_start;
  int i = 0;
  for (; i < (int)ABSL_ARRAYSIZE(tree->sub_tree); ++i) {
    if (tree->sub_tree[i] == nullptr) {
      break;
    }
    // On the right and bottom edges of a frame, the partition tree is populated
    // with 4x4 dummy nodes.
    const bool sub_tree_is_dummy =
        (tree->sub_tree[i]->mi_col == 0 && tree->sub_tree[i]->mi_row == 0 &&
         tree->sub_tree[i]->bsize == BLOCK_4X4 && i > 0);
    if (sub_tree_is_dummy) {
      continue;
    }
    auto *child = partition->add_children();
    child_coding_unit_start =
        PopulatePartitionTree(frame_data, sb_data, child, sb, tree->sub_tree[i],
                              part_type, coeff_idx, child_coding_unit_start);
  }
  coding_unit_end = child_coding_unit_start;
  // No children, so this partition has a coding_unit
  if (i == 0) {
    CodingUnit *cu;
    if (part_type == PartitionType::kChromaOnly) {
      cu = sb->add_coding_units_chroma();
    } else {
      cu = sb->add_coding_units_shared();
    }
    GetCodingUnit(cu, frame_data, sb_data, tree->mi_row, tree->mi_col,
                  part_type, coeff_idx);
    coding_unit_end += 1;
    partition->set_is_leaf_node(true);
  }
  partition->mutable_coding_unit_range()->set_end(coding_unit_end);
  return coding_unit_end;
}

template <typename T>
void PopulateEnumMapping(google::protobuf::Map<int, std::string> *map,
                         const T &enum_names) {
  for (const auto &[enum_value, enum_name] : enum_names) {
    map->insert({ enum_value, std::string(enum_name) });
  }
}

void PopulateEnumMappings(EnumMappings *mappings) {
  PopulateEnumMapping(mappings->mutable_transform_type_mapping(), kTxTypeMap);
  PopulateEnumMapping(mappings->mutable_prediction_mode_mapping(),
                      kPredictionModeMap);
  PopulateEnumMapping(mappings->mutable_uv_prediction_mode_mapping(),
                      kUvPredictionModeMap);
  PopulateEnumMapping(mappings->mutable_motion_mode_mapping(), kMotionModeMap);
  PopulateEnumMapping(mappings->mutable_transform_size_mapping(), kTxSizeMap);
  PopulateEnumMapping(mappings->mutable_block_size_mapping(), kBlockSizeMap);
  PopulateEnumMapping(mappings->mutable_partition_type_mapping(),
                      kPartitionTypeMap);
  PopulateEnumMapping(mappings->mutable_interpolation_filter_mapping(),
                      kInterpFilterMap);
  PopulateEnumMapping(mappings->mutable_entropy_coding_mode_mapping(),
                      kSymbolCodingModeMap);
  PopulateEnumMapping(mappings->mutable_frame_type_mapping(), kFrameTypeMap);
  PopulateEnumMapping(mappings->mutable_tip_mode_mapping(), kTipFrameModeMap);
  PopulateEnumMapping(mappings->mutable_motion_vector_precision_mapping(),
                      kMvSubpelPrecisionMap);
}

bool BlockContains(Position pos, BlockSize size, int x, int y) {
  return pos.x() <= x && x < pos.x() + size.width() && pos.y() <= y &&
         y < pos.y() + size.height();
}

// TODO(comc): Update transform symbol range for transform units.
// Currently accounting context isn't updated to that level of granularity.
void UpdateSymbolRangesCodingUnit(CodingUnit *coding_unit, int index,
                                  int symbol_x, int symbol_y) {
  if (!BlockContains(coding_unit->position(), coding_unit->size(), symbol_x,
                     symbol_y)) {
    return;
  }
  if (!coding_unit->has_symbol_range()) {
    coding_unit->mutable_symbol_range()->set_start(index);
  }
  coding_unit->mutable_symbol_range()->set_end(index + 1);
}

void UpdateSymbolRangesPartition(Superblock *sb, Partition *part, int index,
                                 int symbol_x, int symbol_y, bool is_chroma) {
  if (!BlockContains(part->position(), part->size(), symbol_x, symbol_y)) {
    return;
  }
  if (!part->has_symbol_range()) {
    part->mutable_symbol_range()->set_start(index);
  }
  part->mutable_symbol_range()->set_end(index + 1);

  for (auto &child : *part->mutable_children()) {
    UpdateSymbolRangesPartition(sb, &child, index, symbol_x, symbol_y,
                                is_chroma);
  }

  if (part->is_leaf_node()) {
    uint32_t cu_index = part->coding_unit_range().start();
    CodingUnit *cu = is_chroma ? sb->mutable_coding_units_chroma(cu_index)
                               : sb->mutable_coding_units_shared(cu_index);
    UpdateSymbolRangesCodingUnit(cu, index, symbol_x, symbol_y);
  }
}

void UpdateSymbolRangesSb(Superblock *sb, int index, int symbol_x, int symbol_y,
                          bool is_chroma) {
  auto *part = is_chroma ? sb->mutable_chroma_partition_tree()
                         : sb->mutable_luma_partition_tree();
  UpdateSymbolRangesPartition(sb, part, index, symbol_x, symbol_y, is_chroma);
}

void InspectSuperblock(void *pbi, void *data) {
  auto *ctx = static_cast<ExtractProtoContext *>(data);
  ifd_inspect_superblock(&ctx->frame_data, pbi);
}

void WriteProto(const ExtractProtoContext *ctx, const Frame &frame) {
  auto &frame_params = frame.frame_params();
  std::string_view file_ext =
      ctx->output_config.output_as_text_proto ? kTextProtoExt : kBinaryProtoExt;
  std::string file_prefix = std::string(ctx->output_config.output_prefix);
  if (file_prefix.empty()) {
    file_prefix = ctx->stream_path.stem();
  }
  std::string file_name = absl::StrFormat(
      "%s_frame_%04d.%s", file_prefix, frame_params.decode_index(), file_ext);
  std::filesystem::path output_path =
      ctx->output_config.output_folder / file_name;

  if (ctx->output_config.output_as_text_proto) {
    std::string text_proto;
    CHECK(google::protobuf::TextFormat::PrintToString(frame, &text_proto));
    std::ofstream output_file(output_path);
    output_file << text_proto;
    if (output_file.fail()) {
      LOG(QFATAL) << "Failed to write proto file: " << output_path;
    }
  } else {
    std::ofstream output_file(output_path, std::ofstream::binary);
    frame.SerializeToOstream(&output_file);
    if (output_file.fail()) {
      LOG(QFATAL) << "Failed to write proto file: " << output_path;
    }
  }
  if (ctx->show_progress) {
    std::cout << absl::StrFormat("Wrote: %s\n", output_path) << std::flush;
  }
  if (ctx->output_config.limit != -1 &&
      ctx->decode_count >= ctx->output_config.limit) {
    exit(EXIT_SUCCESS);
  }
}

// See:
// https://gitlab.com/AOMediaCodec/avm/-/blob/4664aa8a08dce15e914093d2c85bcc25d2c8026f/av1/decoder/decodeframe.c#L6899
const int kMaxLagInFrames = 35;
int GetAbsoluteDisplayIndex(ExtractProtoContext *ctx, int raw_display_index,
                            FRAME_TYPE frame_type) {
  int display_index = ctx->display_index_offset + raw_display_index;
  if (ctx->max_display_index - display_index >= kMaxLagInFrames ||
      frame_type == KEY_FRAME) {
    ctx->display_index_offset = ctx->max_display_index + 1;
    display_index = ctx->display_index_offset + raw_display_index;
  }
  ctx->max_display_index = std::max(ctx->max_display_index, display_index);
  return display_index;
}

absl::Status AdvanceToNextY4mMarker(std::ifstream *orig_yuv_file) {
  size_t pos = orig_yuv_file->tellg();
  if (orig_yuv_file->fail()) {
    return absl::UnknownError("Tell failed on YUV file.");
  }
  std::string y4m_data;
  y4m_data.resize(kY4mReadahead);
  orig_yuv_file->read(y4m_data.data(), kY4mReadahead);
  if (orig_yuv_file->fail()) {
    return absl::UnknownError("Read failed on YUV file.");
  }
  size_t next_marker = y4m_data.find(kY4mFrameMarker);
  if (next_marker == std::string::npos) {
    return absl::UnknownError("Unable to find y4m frame marker.");
  }
  orig_yuv_file->seekg(pos + next_marker + kY4mFrameMarker.size());
  if (orig_yuv_file->fail()) {
    return absl::UnknownError("Seek failed on YUV file.");
  }
  return absl::OkStatus();
}

absl::StatusOr<Y4mHeader> ParseY4mHeader(std::ifstream *orig_yuv_file) {
  auto status = AdvanceToNextY4mMarker(orig_yuv_file);
  if (!status.ok()) {
    return status;
  }
  size_t header_end = orig_yuv_file->tellg();
  header_end -= kY4mFrameMarker.size() + 1;
  std::string y4m_header_raw;
  y4m_header_raw.resize(header_end);
  orig_yuv_file->seekg(0);
  orig_yuv_file->read(y4m_header_raw.data(), header_end);
  auto parts = absl::StrSplit(y4m_header_raw, ' ');
  size_t width = 0;
  size_t height = 0;
  size_t bit_depth = 0;
  for (auto &part : parts) {
    if (part.at(0) == 'W') {
      if (!absl::SimpleAtoi(part.substr(1), &width)) {
        return absl::UnknownError(
            absl::StrCat("Invalid width in Y4M header: ", part));
      }
    } else if (part.at(0) == 'H') {
      if (!absl::SimpleAtoi(part.substr(1), &height)) {
        return absl::UnknownError(
            absl::StrCat("Invalid height in Y4M header: ", part));
      }
    } else if (part.at(0) == 'C') {
      if (!kSupportedColorspaces.contains(part)) {
        return absl::UnknownError(
            absl::StrCat("Unsupported Y4M colorspace: ", part));
      }
      bit_depth = kSupportedColorspaces[part];
    }
  }
  if (width == 0 || height == 0 || bit_depth == 0) {
    return absl::UnknownError("Invalid Y4M header.");
  }
  return Y4mHeader{
    .width = width,
    .height = height,
    .bit_depth = bit_depth,
  };
}

struct YuvFrameSizeInfo {
  size_t frame_width;
  size_t frame_height;
  size_t plane_width[3];
  size_t plane_height[3];
  size_t bytes_per_sample;
  size_t frame_size_bytes;
  size_t plane_size_bytes[3];
  size_t plane_offset_bytes[3];
  size_t plane_stride_bytes[3];
};

YuvFrameSizeInfo GetYuvFrameSizeInfo(int bit_depth, size_t frame_width,
                                     size_t frame_height) {
  const size_t bytes_per_sample = (bit_depth == 8) ? 1 : 2;
  const size_t luma_size_bytes = frame_width * frame_height * bytes_per_sample;
  // Round up in the case of odd frame dimensions:
  const size_t chroma_width = (frame_width + 1) / 2;
  const size_t chroma_height = (frame_height + 1) / 2;
  const size_t chroma_size_bytes =
      chroma_width * chroma_height * bytes_per_sample;
  const size_t frame_size_bytes = luma_size_bytes + 2 * chroma_size_bytes;
  return YuvFrameSizeInfo {
    .frame_width = frame_width,
    .frame_height = frame_height,
    .plane_width = {frame_width, chroma_width, chroma_width},
    .plane_height = {frame_height, chroma_height, chroma_height},
    .bytes_per_sample = bytes_per_sample,
    .frame_size_bytes = frame_size_bytes,
    .plane_size_bytes = {
      luma_size_bytes,
      chroma_size_bytes,
      chroma_size_bytes,
    },
    .plane_offset_bytes = {
      0,
      luma_size_bytes,
      luma_size_bytes + chroma_size_bytes,
    },
    .plane_stride_bytes = {
      frame_width * bytes_per_sample,
      chroma_width * bytes_per_sample,
      chroma_width * bytes_per_sample,
    },
  };
}

absl::Status ReadFrameFromYuvFile(ExtractProtoContext *ctx,
                                  size_t frame_display_index,
                                  const YuvFrameSizeInfo &frame_size_info,
                                  std::string &yuv_frame) {
  if (ctx->is_y4m_file) {
    if (ctx->y4m_header.has_value()) {
      if (ctx->y4m_header.value().width != frame_size_info.frame_width ||
          ctx->y4m_header.value().height != frame_size_info.frame_height) {
        LOG(QFATAL) << absl::StrFormat(
            "Y4M dimensions do not match stream: Stream: %dx%d, Y4M: %dx%d",
            frame_size_info.frame_width, frame_size_info.frame_height,
            ctx->y4m_header.value().width, ctx->y4m_header.value().height);
      }
    }

    ctx->orig_yuv_file->seekg(0);
    if (ctx->orig_yuv_file->fail()) {
      return absl::UnknownError("Seek failed on YUV file.");
    }
    for (size_t i = 0; i <= frame_display_index; i++) {
      CHECK_OK(AdvanceToNextY4mMarker(ctx->orig_yuv_file));
      if (i < frame_display_index) {
        ctx->orig_yuv_file->seekg(frame_size_info.frame_size_bytes,
                                  std::ios_base::cur);
      }
    }
  } else {
    ctx->orig_yuv_file->seekg(frame_size_info.frame_size_bytes *
                              frame_display_index);
    if (ctx->orig_yuv_file->fail()) {
      return absl::UnknownError("Seek failed on YUV file.");
    }
  }

  yuv_frame.resize(frame_size_info.frame_size_bytes);
  ctx->orig_yuv_file->read(yuv_frame.data(), frame_size_info.frame_size_bytes);
  if (ctx->orig_yuv_file->fail()) {
    CHECK_OK(absl::UnknownError("Read failed on YUV file."));
  }

  return absl::OkStatus();
}

void CopyFromYuvFile(PixelBuffer *pixel_buffer, const std::string &orig_yuv,
                     int plane, const YuvFrameSizeInfo &frame_size_info,
                     size_t offset_x, size_t offset_y) {
  for (size_t y = offset_y; y < pixel_buffer->height() + offset_y; y++) {
    for (size_t x = offset_x; x < pixel_buffer->width() + offset_x; x++) {
      uint16_t pixel = 0;
      bool in_bounds = x < frame_size_info.plane_width[plane] &&
                       y < frame_size_info.plane_height[plane];
      if (in_bounds) {
        size_t pixel_offset_bytes =
            y * frame_size_info.plane_stride_bytes[plane] +
            x * frame_size_info.bytes_per_sample;
        size_t index =
            frame_size_info.plane_offset_bytes[plane] + pixel_offset_bytes;
        pixel = (uint8_t)orig_yuv[index];
        if (frame_size_info.bytes_per_sample > 1) {
          pixel |= (uint8_t)orig_yuv[index + 1] << 8;
        }
      }
      pixel_buffer->add_pixels(pixel);
    }
  }
}

void InspectTipFrame(void *pbi, void *data) {
  auto *ctx = static_cast<ExtractProtoContext *>(data);
  AV1Decoder *decoder = (AV1Decoder *)pbi;
  AV1_COMMON *const cm = &decoder->common;
  StreamParams *stream_params = ctx->stream_params;
  Frame frame;
  *frame.mutable_stream_params() = *stream_params;
  const int bit_depth = cm->seq_params.bit_depth;
  if (ctx->orig_yuv_bit_depth == -1) {
    ctx->orig_yuv_bit_depth = bit_depth;
  }
  const size_t frame_width = cm->width;
  const size_t frame_height = cm->height;
  YuvFrameSizeInfo frame_size_info =
      GetYuvFrameSizeInfo(ctx->orig_yuv_bit_depth, frame_width, frame_height);

  auto *frame_params = frame.mutable_frame_params();
  frame_params->set_display_index(GetAbsoluteDisplayIndex(
      ctx, cm->current_frame.frame_number, cm->current_frame.frame_type));
  frame_params->set_raw_display_index(cm->current_frame.frame_number);
  frame_params->set_decode_index(ctx->decode_count++);
  frame_params->set_width(frame_width);
  frame_params->set_height(frame_height);
  frame_params->set_bit_depth(bit_depth);
  frame_params->set_frame_type(cm->current_frame.frame_type);
  frame_params->set_show_frame(cm->show_frame);
  PopulateEnumMappings(frame.mutable_enum_mappings());
  auto *tip_frame = frame.mutable_tip_frame_params();
  tip_frame->set_tip_mode(TIP_FRAME_AS_OUTPUT);
  RefCntBuffer *buf = cm->cur_frame;

  std::string orig_yuv = "";
  if (ctx->orig_yuv_file != nullptr) {
    CHECK_OK(ReadFrameFromYuvFile(ctx, frame_params->display_index(),
                                  frame_size_info, orig_yuv));
  }

  for (int plane = 0; plane < 3; ++plane) {
    int plane_width = buf->buf.crop_widths[plane > 0];
    int plane_height = buf->buf.crop_heights[plane > 0];
    int stride = buf->buf.strides[plane > 0];
    auto *pixels = tip_frame->add_pixel_data();
    pixels->set_plane(plane);
    pixels->mutable_reconstruction()->set_width(plane_width);
    pixels->mutable_reconstruction()->set_height(plane_height);
    pixels->mutable_reconstruction()->set_bit_depth(bit_depth);
    for (int y = 0; y < plane_height; ++y) {
      for (int x = 0; x < plane_width; ++x) {
        int16_t recon_pixel = buf->buf.buffers[plane][y * stride + x];
        pixels->mutable_reconstruction()->add_pixels(recon_pixel);
      }
    }

    if (!orig_yuv.empty()) {
      auto original = pixels->mutable_original();
      original->set_width(plane_width);
      original->set_height(plane_height);
      original->set_bit_depth(ctx->orig_yuv_bit_depth);
      CopyFromYuvFile(original, orig_yuv, plane, frame_size_info, 0, 0);
    }
  }
  WriteProto(ctx, frame);
}

void InspectFrame(void *pbi, void *data) {
  ExtractProtoContext *ctx = static_cast<ExtractProtoContext *>(data);
  insp_frame_data &frame_data = ctx->frame_data;
  ifd_inspect(&frame_data, pbi, 0);
  // Show existing frames just show a reference buffer we've already decoded.
  // There's no information to show.
  if (frame_data.show_existing_frame) {
    return;
  }

  StreamParams *stream_params = ctx->stream_params;
  Frame frame;
  *frame.mutable_stream_params() = *stream_params;
  const size_t frame_width = frame_data.width;
  const size_t frame_height = frame_data.height;
  YuvFrameSizeInfo frame_size_info =
      GetYuvFrameSizeInfo(ctx->orig_yuv_bit_depth, frame_width, frame_height);
  auto *frame_params = frame.mutable_frame_params();
  frame_params->set_display_index(GetAbsoluteDisplayIndex(
      ctx, frame_data.frame_number, frame_data.frame_type));
  frame_params->set_raw_display_index(frame_data.frame_number);
  frame_params->set_decode_index(ctx->decode_count++);
  frame_params->set_show_frame(frame_data.show_frame);
  frame_params->set_base_qindex(frame_data.base_qindex);
  frame_params->set_width(frame_width);
  frame_params->set_height(frame_height);

  auto *tip_frame = frame.mutable_tip_frame_params();
  tip_frame->set_tip_mode(frame_data.tip_frame_mode);
  const int sb_width_mi = mi_size_wide[frame_data.superblock_size];
  const int sb_height_mi = mi_size_high[frame_data.superblock_size];
  const int sb_width_px = sb_width_mi * MI_SIZE;
  const int sb_height_px = sb_height_mi * MI_SIZE;
  frame_params->mutable_superblock_size()->set_width(sb_width_px);
  frame_params->mutable_superblock_size()->set_height(sb_height_px);
  frame_params->mutable_superblock_size()->set_enum_value(
      frame_data.superblock_size);
  frame_params->set_frame_type(frame_data.frame_type);
  int bit_depth = frame_data.bit_depth;
  if (ctx->orig_yuv_bit_depth == -1) {
    ctx->orig_yuv_bit_depth = bit_depth;
  }
  frame_params->set_bit_depth(bit_depth);
  PopulateEnumMappings(frame.mutable_enum_mappings());
  std::string orig_yuv = "";
  if (ctx->orig_yuv_file != nullptr) {
    CHECK_OK(ReadFrameFromYuvFile(ctx, frame_params->display_index(),
                                  frame_size_info, orig_yuv));
  }
  const Accounting *accounting = frame_data.accounting;
  const int num_symbol_types = accounting->syms.dictionary.num_strs;
  auto *symbol_info = frame.mutable_symbol_info();
  for (int i = 0; i < num_symbol_types; i++) {
    SymbolInfo info;
    info.set_source_file(accounting->syms.dictionary.acct_infos[i].c_file);
    info.set_source_line(accounting->syms.dictionary.acct_infos[i].c_line);
    info.set_source_function(accounting->syms.dictionary.acct_infos[i].c_func);
    for (int tag = 0; tag < AOM_ACCOUNTING_MAX_TAGS; tag++) {
      if (accounting->syms.dictionary.acct_infos[i].tags[tag] != nullptr) {
        info.add_tags(accounting->syms.dictionary.acct_infos[i].tags[tag]);
      }
    }
    (*symbol_info)[i] = info;
  }
  int sb_rows = (frame_data.mi_rows + sb_width_mi - 1) / sb_width_mi;
  int sb_cols = (frame_data.mi_cols + sb_height_mi - 1) / sb_height_mi;
  for (int sb_row = 0; sb_row < sb_rows; sb_row++) {
    for (int sb_col = 0; sb_col < sb_cols; sb_col++) {
      auto *sb = frame.add_superblocks();
      int sb_x_offset_px = sb_col * sb_width_px;
      int sb_y_offset_px = sb_row * sb_height_px;
      sb->mutable_position()->set_x(sb_x_offset_px);
      sb->mutable_position()->set_y(sb_y_offset_px);
      sb->mutable_size()->set_width(sb_width_px);
      sb->mutable_size()->set_height(sb_height_px);
      sb->mutable_size()->set_enum_value(frame_data.superblock_size);
      // TODO(comc): Add special case for monochrome
      for (int plane = 0; plane < 3; plane++) {
        int sb_plane_x_offset_px = plane ? sb_x_offset_px / 2 : sb_x_offset_px;
        int sb_plane_y_offset_px = plane ? sb_y_offset_px / 2 : sb_y_offset_px;
        int sb_plane_width_px = plane ? sb_width_px / 2 : sb_width_px;
        int sb_plane_height_px = plane ? sb_height_px / 2 : sb_height_px;
        int remaining_width =
            frame_size_info.plane_width[plane] - sb_plane_x_offset_px;
        int remaining_height =
            frame_size_info.plane_height[plane] - sb_plane_y_offset_px;
        int cropped_sb_plane_width_px =
            std::min(sb_plane_width_px, remaining_width);
        int cropped_sb_plane_height_px =
            std::min(sb_plane_height_px, remaining_height);
        auto *pixels = sb->add_pixel_data();
        pixels->set_plane(plane);
        pixels->mutable_reconstruction()->set_width(sb_plane_width_px);
        pixels->mutable_reconstruction()->set_height(sb_plane_height_px);
        pixels->mutable_reconstruction()->set_bit_depth(frame_data.bit_depth);

        pixels->mutable_pre_filtered()->set_width(sb_plane_width_px);
        pixels->mutable_pre_filtered()->set_height(sb_plane_height_px);
        pixels->mutable_pre_filtered()->set_bit_depth(frame_data.bit_depth);

        pixels->mutable_prediction()->set_width(sb_plane_width_px);
        pixels->mutable_prediction()->set_height(sb_plane_height_px);
        pixels->mutable_prediction()->set_bit_depth(frame_data.bit_depth);
        for (int px_y = 0; px_y < sb_plane_height_px; px_y++) {
          for (int px_x = 0; px_x < sb_plane_width_px; px_x++) {
            int stride = frame_data.recon_frame_buffer.strides[plane > 0];
            int pixel_y = sb_plane_y_offset_px + px_y;
            int pixel_x = sb_plane_x_offset_px + px_x;
            int pixel_offset = pixel_y * stride + pixel_x;
            bool in_bounds = px_y < cropped_sb_plane_height_px &&
                             px_x < cropped_sb_plane_width_px;
            uint16_t recon_pixel =
                in_bounds
                    ? frame_data.recon_frame_buffer.buffers[plane][pixel_offset]
                    : 0;
            pixels->mutable_reconstruction()->add_pixels(recon_pixel);

            uint16_t pre_filtered_pixel =
                in_bounds ? frame_data.prefiltered_frame_buffer
                                .buffers[plane][pixel_offset]
                          : 0;
            pixels->mutable_pre_filtered()->add_pixels(pre_filtered_pixel);

            uint16_t predicted_pixel = in_bounds
                                           ? frame_data.predicted_frame_buffer
                                                 .buffers[plane][pixel_offset]
                                           : 0;
            pixels->mutable_prediction()->add_pixels(predicted_pixel);
          }
        }

        if (!orig_yuv.empty()) {
          auto original = pixels->mutable_original();
          original->set_width(sb_plane_width_px);
          original->set_height(sb_plane_height_px);
          original->set_bit_depth(ctx->orig_yuv_bit_depth);
          CopyFromYuvFile(original, orig_yuv, plane, frame_size_info,
                          sb_plane_x_offset_px, sb_plane_y_offset_px);
        }
      }
      insp_sb_data *sb_data =
          &frame_data.sb_grid[sb_row * frame_data.max_sb_cols + sb_col];
      auto *luma_part = sb->mutable_luma_partition_tree();
      auto *chroma_part = sb->mutable_chroma_partition_tree();
      int coeff_idx[3] = { 0, 0, 0 };
      sb->set_has_separate_chroma_partition_tree(
          sb_data->has_separate_chroma_partition_tree);

      auto primary_part_type = sb_data->has_separate_chroma_partition_tree
                                   ? PartitionType::kLumaOnly
                                   : PartitionType::kShared;
      PopulatePartitionTree(&frame_data, sb_data, luma_part, sb,
                            sb_data->partition_tree_luma, primary_part_type,
                            coeff_idx, 0);
      if (sb_data->has_separate_chroma_partition_tree) {
        PopulatePartitionTree(&frame_data, sb_data, chroma_part, sb,
                              sb_data->partition_tree_chroma,
                              PartitionType::kChromaOnly, coeff_idx, 0);
      }
    }
  }

  int prev_sb_i = -1;
  int relative_index = 0;
  const int num_syms = accounting->syms.num_syms;
  for (int i = 0; i < num_syms; i++) {
    auto *symbol = &accounting->syms.syms[i];
    if (symbol->context.x == -1 || symbol->context.y == -1) {
      symbol->context.x = 0;
      symbol->context.y = 0;
    }
    bool is_chroma = symbol->context.tree_type == CHROMA_PART;
    int sb_col = symbol->context.x / sb_width_mi;
    int sb_row = symbol->context.y / sb_height_mi;
    int sb_i = sb_row * sb_cols + sb_col;
    if (sb_i != prev_sb_i) {
      relative_index = 0;
    }
    auto *sb = frame.mutable_superblocks(sb_i);
    UpdateSymbolRangesSb(sb, relative_index, symbol->context.x * MI_SIZE,
                         symbol->context.y * MI_SIZE, is_chroma);
    auto *sym = sb->add_symbols();
    sym->set_bits(symbol->bits / (float)(1 << AOM_ACCT_BITRES));
    sym->set_value(symbol->value);
    sym->set_info_id(symbol->id);
    prev_sb_i = sb_i;
    relative_index += 1;
  }
  WriteProto(ctx, frame);
}

void SetupInspectCallbacks(ExtractProtoContext *ctx) {
  aom_inspect_init ii;
  ii.inspect_cb = InspectFrame;
  ii.inspect_sb_cb = InspectSuperblock;
  ii.inspect_tip_cb = InspectTipFrame;
  ii.inspect_ctx = static_cast<void *>(ctx);
  aom_codec_control(&ctx->codec, AV1_SET_INSPECTION_CALLBACK, &ii);
}

void PopulateStreamPath(ExtractProtoContext *ctx) {
  auto absolute_path = std::filesystem::absolute(ctx->stream_path);
  std::filesystem::path relative_path;
  if (ctx->preserve_stream_path_depth == -1) {
    relative_path = absolute_path;
  } else {
    int index = std::distance(absolute_path.begin(), absolute_path.end());
    for (const auto &part : absolute_path) {
      index -= 1;
      if (index <= ctx->preserve_stream_path_depth) {
        relative_path /= part;
      }
    }
  }
  ctx->stream_params->set_stream_path(relative_path.make_preferred());
}

absl::Status OpenStream(ExtractProtoContext *ctx) {
  ctx->reader = aom_video_reader_open(ctx->stream_path.c_str());
  if (ctx->reader == nullptr) {
    return absl::NotFoundError(
        absl::StrFormat("Failed to open %s for reading.", ctx->stream_path));
  }
  ctx->info = aom_video_reader_get_info(ctx->reader);
  aom_codec_iface_t *decoder =
      get_aom_decoder_by_fourcc(ctx->info->codec_fourcc);
  if (decoder == nullptr) {
    return absl::UnimplementedError("Unknown input codec.");
  }
  LOG(INFO) << "Using " << aom_codec_iface_name(decoder);
  ctx->stream_params->set_avm_version(aom_codec_iface_name(decoder));
  ctx->stream_params->set_stream_name(ctx->stream_path.filename());
  PopulateStreamPath(ctx);
  if (aom_codec_dec_init(&ctx->codec, decoder, nullptr, 0)) {
    return absl::InternalError("Failed to initialize decoder.");
  }
  ifd_init(&ctx->frame_data, ctx->info->frame_width, ctx->info->frame_height);
  SetupInspectCallbacks(ctx);
  return absl::OkStatus();
}

// Note: TIP can mean the number of decoded frames is significantly less than
// the number of displayed frames.
absl::Status ReadFrames(ExtractProtoContext *ctx) {
  const uint8_t *frame;
  size_t frame_size = 0;
  int frame_count = 0;
  while (true) {
    if (!aom_video_reader_read_frame(ctx->reader)) {
      return absl::OkStatus();
    }
    frame = aom_video_reader_get_frame(ctx->reader, &frame_size);
    // TODO(comc): Check handling of non-display frames; user_priv=nullptr
    // bypasses the custom decoder_inspect logic in av1_dx_iface.c.
    if (aom_codec_decode(&ctx->codec, frame, (unsigned int)frame_size,
                         nullptr) != AOM_CODEC_OK) {
      return absl::InternalError(absl::StrFormat("Failed to decode frame: %s",
                                                 aom_codec_error(&ctx->codec)));
    }

    bool got_any_frames = false;
    av1_ref_frame ref_dec;
    ref_dec.idx = -1;

    aom_image_t *img = nullptr;
    (void)img;
    // ref_dec.idx is the index to the reference buffer idx to AV1_GET_REFERENCE
    // if its -1 the decoder didn't update any reference buffer and the only
    // way to see the frame is aom_codec_get_frame.
    if (ref_dec.idx == -1) {
      aom_codec_iter_t iter = nullptr;
      while ((img = aom_codec_get_frame(&ctx->codec, &iter)) != nullptr) {
        ++frame_count;
      }
      got_any_frames = true;
    } else if (!aom_codec_control(&ctx->codec, AV1_GET_REFERENCE, &ref_dec)) {
      img = &ref_dec.img;
      ++frame_count;
      got_any_frames = true;
    }
    if (!got_any_frames) {
      return absl::InternalError("No frames decoded.");
    }
  }
}
}  // namespace

void ValidateFlagPath(const absl::Flag<std::string> &flag) {
  auto val = absl::GetFlag(flag);
  QCHECK(!val.empty()) << absl::StrFormat("--%s must be set.", flag.Name());
  QCHECK(std::filesystem::exists(val))
      << absl::StrFormat("Path %s does not exist.", val);
}

int main(int argc, char **argv) {
  absl::ParseCommandLine(argc, argv);
  absl::InitializeLog();
  absl::SetStderrThreshold(absl::LogSeverity::kInfo);
  ValidateFlagPath(FLAGS_stream);
  ValidateFlagPath(FLAGS_output_folder);
  std::string yuv_contents = "";
  const std::filesystem::path stream_path = absl::GetFlag(FLAGS_stream);
  const std::filesystem::path output_folder =
      absl::GetFlag(FLAGS_output_folder);
  const std::filesystem::path orig_yuv_path = absl::GetFlag(FLAGS_orig_yuv);
  std::ifstream orig_yuv_file;
  bool is_y4m_file = false;
  bool have_orig_yuv = !orig_yuv_path.empty();
  if (have_orig_yuv) {
    auto ext = orig_yuv_path.extension();
    if (ext == ".y4m") {
      is_y4m_file = true;
    } else if (ext != ".yuv") {
      LOG(QFATAL)
          << "Invalid YUV file extension (expected either .y4m or .yuv): "
          << orig_yuv_path;
    }
    orig_yuv_file.open(orig_yuv_path, std::ifstream::binary);
    if (orig_yuv_file.fail()) {
      LOG(QFATAL) << "Failed to open YUV file.";
    }
  }

  StreamParams params;
  params.mutable_stream_name()->assign(stream_path);

  OutputConfig output_config = {
    .output_folder = output_folder,
    .output_prefix = absl::GetFlag(FLAGS_output_prefix),
    .output_as_text_proto = absl::GetFlag(FLAGS_output_as_text),
    .limit = absl::GetFlag(FLAGS_limit),
  };

  int orig_yuv_bit_depth = absl::GetFlag(FLAGS_orig_yuv_bit_depth);
  std::optional<Y4mHeader> y4m_header = std::nullopt;
  if (is_y4m_file && !absl::GetFlag(FLAGS_ignore_y4m_header)) {
    auto y4m_reader_status = ParseY4mHeader(&orig_yuv_file);
    if (y4m_reader_status.ok()) {
      y4m_header = y4m_reader_status.value();
      if (orig_yuv_bit_depth != -1) {
        LOG(WARNING) << absl::StrFormat(
            "orig_yuv_bit_depth was manually specified, but overridden by the "
            "y4m header.");
      }
      orig_yuv_bit_depth = y4m_header.value().bit_depth;
    } else {
      LOG(QFATAL) << y4m_reader_status.status();
    }
  }
  ExtractProtoContext ctx = {
    .frame_data = {},
    .codec = {},
    .reader = nullptr,
    .info = nullptr,
    .stream_params = &params,
    .orig_yuv_file = have_orig_yuv ? &orig_yuv_file : nullptr,
    .orig_yuv_bit_depth = orig_yuv_bit_depth,
    .stream_path = stream_path,
    .preserve_stream_path_depth =
        absl::GetFlag(FLAGS_preserve_stream_path_depth),
    .decode_count = 0,
    .is_y4m_file = is_y4m_file,
    .y4m_header = y4m_header,
    .output_config = output_config,
    .show_progress = absl::GetFlag(FLAGS_show_progress),
    .display_index_offset = 0,
    .max_display_index = -1,
  };
  CHECK_OK(OpenStream(&ctx));

  // TODO(comc): These are sometimes 0 after the OpenStream call.
  params.set_width(ctx.info->frame_width);
  params.set_height(ctx.info->frame_height);
  for (const auto &enc_arg : absl::GetFlag(FLAGS_encoder_args)) {
    std::pair<std::string, std::string> pair =
        absl::StrSplit(enc_arg, absl::MaxSplits('=', 1));
    params.mutable_encoder_args()->insert({ pair.first, pair.second });
  }

  CHECK_OK(ReadFrames(&ctx));
  return EXIT_SUCCESS;
}

// Linking against AVM requires this symbol to be defined.
void usage_exit(void) {
  LOG(QFATAL) << "Usage: extract_proto src_filename <options>";
}