#!/usr/bin/env python3
# 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/.
##
"""Creates enum value to name mapping header for AVM.

Walks through the enum definitions in av1/common/enums.h and creates a header that
maps individual enum variants to their (string) names.
"""

import argparse
from collections.abc import Sequence
import re

parser = argparse.ArgumentParser()
parser.add_argument("--enums_header", action="append", required=True)
parser.add_argument("--output", required=True)

# AVM enum definitions start with an optional typedef, the enum keyword, and
# then optionally `ATTRIBUTE_PACKED`.
ENUM_START_REGEX = re.compile(r"(?:typedef )?enum (?:ATTRIBUTE_PACKED )?\{")
# AVM enum definitions optionally end with the macro UENUM1BYTE or SENUM1BYTE
# and the enum name.
ENUM_END_REGEX = re.compile(r"} ([US]ENUM[12]BYTE\()?(\w+)(?(1)\)|);")
# Some enum variants are defined based on another variant; we don't care about
# these ones, and only want to handle the ones with unique values. These variants
# typically appear as actual symbols in the stream; the others are defined just
# for convenience to mark the start and endpoints of meaningful groups of enums,
# e.g. INTRA_MODE_START / INTRA_MODE_END are used to mark which prediction modes
# are intra (rather than inter).
ENUM_VARIANT_REGEX = re.compile(r"(\w+)(=(\d+))?,")
# String template for each enum mapping.
ENUM_MAP_TEMPLATE = """static absl::flat_hash_map<int, std::string_view> k{name_camel}Map = {{"""
# These preprocessor directives are replaced with dummy conditions to simplify
# handling of their corresponding #endif. #ifdef/#ifndef are only used for
# include guards and #ifdef __cplusplus, which we can ignore for the generated
# header.
SKIP_DIRECTIVES = ("#ifdef", "#ifndef")
# These preprocessor directives are relevant for enum definitions, so pass them
# through to the generated header as-is. Ignore all other preprocessor
# directives.
PASSTHROUGH_DIRECTIVES = ("#if ", "#elif", "#else", "#endif")

# Template for the generated header.
ENUM_HEADER_TEMPLATE = """
/*
 * 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/.
 */
// Auto-generated by generate_enum_mappings.py.
#ifndef AOM_TOOLS_EXTRACT_PROTO_ENUM_MAPPINGS_H_
#define AOM_TOOLS_EXTRACT_PROTO_ENUM_MAPPINGS_H_
#include <string_view>

#include "av1/common/av1_common_int.h"
#include "av1/common/blockd.h"
#include "av1/common/enums.h"
#include "av1/common/filter.h"
#include "av1/common/mv.h"
#include "av1/decoder/accounting.h"
#include "config/aom_config.h"
#include "absl/container/flat_hash_map.h"

{enum_definitions}

#endif  // AOM_TOOLS_EXTRACT_PROTO_ENUM_MAPPINGS_H_
"""


def snake_to_camel_case(s: str) -> str:
  # Some enum names are already camel case, so return as-is
  if "_" not in s or not s.isupper():
    return s
  return "".join(c.capitalize() for c in s.lower().split("_"))


def strip_comments(s: str) -> str:
  # Note: Assumes no nested comments
  s = re.sub(r"\/\*.+?\*\/", "", s, flags=re.DOTALL)
  s = re.sub(r"//.*", "", s)
  return s


def remove_backslashes(s: str) -> str:
  s = re.sub(r"\\\n", "", s, flags=re.DOTALL)
  return s

# TODO(comc) Unit test this, and its helper functions.
def create_enum_mapping(input_header_paths: Sequence[str], output_header_path: str):
  enum_definitions = []
  for input_header_path in input_header_paths:
    with open(input_header_path, "r") as f:
      header = f.read()
    header = strip_comments(header)
    header = remove_backslashes(header)
    output_lines = []
    current_enum_index = None
    for line in header.split("\n"):
      line = line.strip()
      if any(line.startswith(s) for s in PASSTHROUGH_DIRECTIVES):
        output_lines.append(line)
      elif any(line.startswith(s) for s in SKIP_DIRECTIVES):
        # To simplify parsing, when we encounter #ifdef / #ifndef, just replace it
        # with `#if 1` so the corresponding #endif still works.
        output_lines.append("#if 1")
      elif re.fullmatch(ENUM_START_REGEX, line):
        current_enum_index = len(output_lines)
        output_lines.append(ENUM_MAP_TEMPLATE)
      elif current_enum_index is not None:
        if variant := re.fullmatch(ENUM_VARIANT_REGEX, line.replace(" ", "")):
          name = variant.group(1)
          output_lines.append(f'    {{{name}, "{name}"}},')
        elif enum_name := re.fullmatch(ENUM_END_REGEX, line):
          output_lines[current_enum_index] = ENUM_MAP_TEMPLATE.format(
              name_camel=snake_to_camel_case(enum_name.group(2))
          )
          current_enum_index = None
          output_lines.append("};")
    enum_definitions.append("\n".join(output_lines))
  enum_header = ENUM_HEADER_TEMPLATE.format(
      enum_definitions="\n".join(enum_definitions))
  with open(output_header_path, "w") as f:
    f.write(enum_header)


def main() -> None:
  args = parser.parse_args()
  create_enum_mapping(args.enums_header, args.output)


if __name__ == "__main__":
  main()