#!/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", 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 have a defined value; we don't care about these ones, and
# only want to handle the ones with auto-incremented values. The
# auto-incremented variants actually appear as 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+),")
# 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/enums.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


# TODO(comc) Unit test this, and its helper functions.
def create_enum_mapping(input_header_path: str, output_header_path: str):
  with open(input_header_path, "r") as f:
    header = f.read()
  header = strip_comments(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.strip()):
        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 = "\n".join(output_lines)
  enum_header = ENUM_HEADER_TEMPLATE.format(enum_definitions=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()