# Copyright (C) 2022 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List
from typing import Optional

from python.generators.trace_processor_table.public import Alias
from python.generators.trace_processor_table.public import ColumnFlag
from python.generators.trace_processor_table.public import CppAccess
from python.generators.trace_processor_table.public import CppAccessDuration
from python.generators.trace_processor_table.public import Purpose
from python.generators.trace_processor_table.public import SqlAccess
from python.generators.trace_processor_table.util import ParsedTable
from python.generators.trace_processor_table.util import ParsedColumn
from python.generators.trace_processor_table.util import parse_type


class ColumnSerializer:

  def __init__(self, table: ParsedTable, column: ParsedColumn, col_index: int):
    self.col_index = col_index
    self.parsed_col = column
    self.col = self.parsed_col.column
    self.name = self.col.name
    self.flags = self.col.flags

    parsed_type = parse_type(table.table, self.col.type)

    self.cpp_type_non_optional = parsed_type.cpp_type()
    self.cpp_type_with_optionality = parsed_type.cpp_type_with_optionality()
    self.is_optional = parsed_type.is_optional
    self.is_string = self.cpp_type_non_optional == 'StringPool::Id'
    self.is_id_type = parsed_type.id_table
    self.is_no_transform_id = parsed_type.is_no_transform_id()

    self.is_implicit_id = self.parsed_col.is_implicit_id
    self.table_purpose = table.table.purpose

  def row_field(self) -> Optional[str]:
    if self.is_implicit_id:
      return None
    return f'    {self.cpp_type_with_optionality} {self.name};'

  def row_insert_arg_value(self) -> Optional[str]:
    """Generates the C++ code to access this column's value from a Row object,
    for use in the dataframe::InsertUnchecked call."""
    if self.is_implicit_id:
      return f'std::monostate()'
    if self.is_optional and self.is_id_type and not self.is_no_transform_id:
      return f'row.{self.name} ? std::make_optional(row.{self.name}->value) : std::nullopt'
    if self.is_optional and self.is_string:
      return f'row.{self.name}'
    if self.is_string:
      return f'std::make_optional(row.{self.name})'
    if self.is_id_type and not self.is_no_transform_id:
      return f'row.{self.name}.value'
    return f'row.{self.name}'

  def cursor_getter(self) -> Optional[str]:
    if self.col.cpp_access == CppAccess.NONE:
      return ''

    if self.col.cpp_access_duration == CppAccessDuration.PRE_FINALIZATION_ONLY:
      dcheck = 'PERFETTO_DCHECK(!dataframe_->finalized());'
    else:
      dcheck = ''

    if self.is_optional and self.is_id_type:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        auto res = cursor_.GetCellUnchecked<ColumnIndex::{self.name}>(kSpec);
        return res ? std::make_optional({self.cpp_type_non_optional}{{*res}}) : std::nullopt;
      }}'''
    if self.is_optional and self.is_string:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        return cursor_.GetCellUnchecked<ColumnIndex::{self.name}>(kSpec);
      }}'''
    if self.is_string:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        auto res = cursor_.GetCellUnchecked<ColumnIndex::{self.name}>(kSpec);
        return res ? *res : StringPool::Id::Null();
      }}'''
    if self.is_id_type:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        return {self.cpp_type_non_optional}{{cursor_.GetCellUnchecked<ColumnIndex::{self.name}>(kSpec)}};
      }}'''
    return f'''
    {self.cpp_type_with_optionality} {self.name}() const {{
      {dcheck}
      return cursor_.GetCellUnchecked<ColumnIndex::{self.name}>(kSpec);
    }}'''

  def cursor_setter(self) -> Optional[str]:
    if self.col.cpp_access == CppAccess.NONE or self.col.cpp_access == CppAccess.READ:
      return ''
    if self.col.cpp_access_duration == CppAccessDuration.PRE_FINALIZATION_ONLY:
      dcheck = 'PERFETTO_DCHECK(!dataframe_->finalized());'
    else:
      dcheck = ''
    if self.is_optional and self.is_id_type and not self.is_no_transform_id:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        auto res_value = res ? std::make_optional(res->value) : std::nullopt;
        cursor_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, res_value);
      }}'''
    if self.is_optional and self.is_string:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        cursor_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, res);
    }}'''
    if self.is_string:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        cursor_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, std::make_optional(res));
    }}'''
    if self.is_id_type and not self.is_no_transform_id:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        cursor_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, res.value);
      }}'''
    return f'''
    void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
      cursor_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, res);
    }}'''

  def row_reference_getter(self) -> str:
    if self.col.cpp_access == CppAccess.NONE:
      return ''
    if self.col.cpp_access_duration == CppAccessDuration.PRE_FINALIZATION_ONLY:
      dcheck = 'PERFETTO_DCHECK(!table_->dataframe_.finalized());'
    else:
      dcheck = ''
    if self.is_optional and self.is_id_type:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        auto res = table_->dataframe_.template GetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_);
        return res ? std::make_optional({self.cpp_type_non_optional}{{*res}}) : std::nullopt;
      }}'''
    if self.is_optional and self.is_string:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        return table_->dataframe_.template GetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_);
      }}'''
    if self.is_string:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        auto res = table_->dataframe_.template GetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_);
        return res ? *res : StringPool::Id::Null();
      }}'''
    if self.is_id_type:
      return f'''
      {self.cpp_type_with_optionality} {self.name}() const {{
        {dcheck}
        return {self.cpp_type_non_optional}{{table_->dataframe_.template GetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_)}};
      }}'''
    return f'''
    {self.cpp_type_with_optionality} {self.name}() const {{
      {dcheck}
      return table_->dataframe_.template GetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_);
    }}'''

  def row_reference_setter(self) -> Optional[str]:
    if self.col.cpp_access == CppAccess.NONE or self.col.cpp_access == CppAccess.READ:
      return ''
    if self.col.cpp_access_duration == CppAccessDuration.PRE_FINALIZATION_ONLY:
      dcheck = 'PERFETTO_DCHECK(!table_->dataframe_.finalized());'
    else:
      dcheck = ''
    if self.is_optional and self.is_id_type and not self.is_no_transform_id:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        auto res_value = res ? std::make_optional(res->value) : std::nullopt;
        table_->dataframe_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_, res_value);
      }}'''
    if self.is_optional and self.is_string:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        table_->dataframe_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_, res);
    }}'''
    if self.is_string:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        table_->dataframe_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_, std::make_optional(res));
    }}'''
    if self.is_id_type and not self.is_no_transform_id:
      return f'''
      void set_{self.name}({self.cpp_type_with_optionality} res) {{
        {dcheck}
        table_->dataframe_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_, res.value);
      }}'''
    return f'''
    void set_{self.name}({self.cpp_type_with_optionality} res) {{
      {dcheck}
      table_->dataframe_.SetCellUnchecked<ColumnIndex::{self.name}>(kSpec, row_, res);
    }}'''

  def colindex_member(self) -> str:
    return f'static constexpr uint32_t {self.name} = {self.col_index}'

  def typespec_column_name_literal(self) -> str:
    return f'"{self.name}"'

  def _get_storage_type_tag(self) -> str:
    if self.is_implicit_id:
      return 'dataframe::Id'
    if self.is_string:
      return 'dataframe::String'
    if self.cpp_type_non_optional == 'uint32_t' or self.is_id_type:
      return 'dataframe::Uint32'
    if self.cpp_type_non_optional == 'int32_t':
      return 'dataframe::Int32'
    if self.cpp_type_non_optional == 'int64_t':
      return 'dataframe::Int64'
    if self.cpp_type_non_optional == 'double':
      return 'dataframe::Double'
    raise ValueError(
        f"Unknown cpp_type_non_optional '{self.cpp_type_non_optional}' for Dataframe StorageType for column '{self.name}'"
    )

  def _get_nullability_tag(self) -> str:
    if self.is_optional or self.is_string:
      # StaticTableFunction outputs are read by random access (GetCell), which
      # the scan-only SparseNull layout does not support. Force DenseNull for
      # all nullable columns regardless of cpp_access.
      if self.table_purpose == Purpose.STATIC_TABLE_FUNCTION:
        return 'dataframe::DenseNull'
      if self.col.sql_access == SqlAccess.HIGH_PERF:
        return 'dataframe::DenseNull'
      if self.col.cpp_access == CppAccess.NONE:
        return 'dataframe::SparseNull'
      if self.col.cpp_access == CppAccess.READ_AND_HIGH_PERF_WRITE:
        return 'dataframe::DenseNull'
      assert self.col.cpp_access in (
          CppAccess.READ,
          CppAccess.READ_AND_LOW_PERF_WRITE,
      )
      # TODO(lalitm): we force this to dense for strings because in the
      # pre-dataframe days, string nullability was just represented by
      # StringPool::Id::Null(). There are hidden dependencies everywhere in the
      # codebase that expect O(1) writes to string columns. Fixing this is
      # non-trivial, so we just keep it dense for now.
      if self.is_string and self.col.cpp_access == CppAccess.READ_AND_LOW_PERF_WRITE:
        return 'dataframe::DenseNull'
      if self.col.cpp_access_duration == CppAccessDuration.PRE_FINALIZATION_ONLY:
        return 'dataframe::SparseNullWithPopcountUntilFinalization'
      assert self.col.cpp_access_duration == CppAccessDuration.POST_FINALIZATION
      return 'dataframe::SparseNullWithPopcountAlways'
    return 'dataframe::NonNull'

  def _get_sort_state_tag(self) -> str:
    if ColumnFlag.SORTED in self.flags:
      if self.is_implicit_id:
        return 'dataframe::IdSorted'
      if ColumnFlag.SET_ID in self.flags:
        return 'dataframe::SetIdSorted'
      return 'dataframe::Sorted'
    return 'dataframe::Unsorted'

  def _get_duplicate_state_tag(self) -> str:
    if self.is_implicit_id:
      return 'dataframe::NoDuplicates'
    return 'dataframe::HasDuplicates'

  def typespec_typed_column_spec_expr(self) -> str:
    storage_tag = self._get_storage_type_tag()
    null_tag = self._get_nullability_tag()
    sort_tag = self._get_sort_state_tag()
    dupe_tag = self._get_duplicate_state_tag()
    return (
        f"dataframe::CreateTypedColumnSpec("
        f"{storage_tag}{{}}, {null_tag}{{}}, {sort_tag}{{}}, {dupe_tag}{{}})")

  def column_spec_initializer(self) -> str:
    """Generate ColumnSpec initializer for non-templated array."""
    storage_tag = self._get_storage_type_tag()
    null_tag = self._get_nullability_tag()
    sort_tag = self._get_sort_state_tag()
    dupe_tag = self._get_duplicate_state_tag()
    return f"dataframe::ColumnSpec{{{storage_tag}{{}}, {null_tag}{{}}, {sort_tag}{{}}, {dupe_tag}{{}}}}"


class TableSerializer(object):

  def __init__(self, parsed: ParsedTable):
    self.table = parsed.table
    self.table_name = parsed.table.class_name
    self.column_serializers: List[ColumnSerializer] = []

    # Filter out aliases first, then assign col_index
    temp_serializers = []
    for c_parsed in parsed.columns:
      if isinstance(c_parsed.column.type, Alias):
        continue
      # col_index will be assigned when rebuilding self.column_serializers
      temp_serializers.append(ColumnSerializer(parsed, c_parsed, 0))

    self.column_serializers = [
        ColumnSerializer(parsed, cs.parsed_col, i)
        for i, cs in enumerate(temp_serializers)
    ]

  def foreach_col(self, serialize_fn, delimiter='\n') -> str:
    lines = []
    for c_ser in self.column_serializers:
      serialized = serialize_fn(c_ser)
      if serialized:
        lines.append(serialized.lstrip('\n').rstrip())
    return delimiter.join(lines).strip()

  def serialize_id_type(self) -> str:
    """Generate the Id type at namespace scope."""
    return f'''struct {self.table_name}_Id : BaseId {{
  {self.table_name}_Id() = default;
  explicit constexpr {self.table_name}_Id(uint32_t _value) : BaseId(_value) {{}}

  bool operator==(const {self.table_name}_Id& other) const {{
    return value == other.value;
  }}
}};'''

  def serialize_static_specs(self) -> str:
    """Generate non-templated column names and specs arrays."""
    col_names = self.foreach_col(
        ColumnSerializer.typespec_column_name_literal, delimiter=', ')
    col_specs = self.foreach_col(
        ColumnSerializer.column_spec_initializer, delimiter=',\n    ')
    col_count = len(self.column_serializers)

    return f'''
inline constexpr const char* k{self.table_name}ColumnNames[] = {{{col_names}}};
inline constexpr dataframe::ColumnSpec k{self.table_name}ColumnSpecs[] = {{
    {col_specs}}};
inline constexpr uint32_t k{self.table_name}ColumnCount = {col_count};
'''

  def row_struct(self) -> str:
    fields = []
    constructor_params_list = []
    constructor_inits_list = []

    for c_ser in self.column_serializers:
      if not c_ser.is_implicit_id:
        fields.append(c_ser.row_field())
        constructor_params_list.append(
            f'{c_ser.cpp_type_with_optionality} _{c_ser.name} = {{}}')
        constructor_inits_list.append(f'{c_ser.name}(std::move(_{c_ser.name}))')

    fields_str = "\n".join(filter(None, fields))
    params_str = ', '.join(constructor_params_list)
    inits_str = (
        ': ' +
        ', '.join(constructor_inits_list)) if constructor_inits_list else ''

    return f'''
  struct Row {{
    Row({params_str}) {inits_str} {{}}

    bool operator==(const Row& other) const {{
      return {' && '.join([f'std::equal_to<>()({c_ser.name}, other.{c_ser.name})' for c_ser in self.column_serializers if not c_ser.is_implicit_id])};
    }}

    {fields_str}
  }};'''

  def serialize(self) -> str:
    col_names = self.foreach_col(
        ColumnSerializer.typespec_column_name_literal, delimiter=',')
    col_specs = self.foreach_col(
        ColumnSerializer.typespec_typed_column_spec_expr, delimiter=',\n    ')
    col_index = self.foreach_col(
        ColumnSerializer.colindex_member, delimiter=';\n    ')
    row_ref_getter = self.foreach_col(
        ColumnSerializer.row_reference_getter, delimiter='\n    ')
    row_ref_setter = self.foreach_col(
        ColumnSerializer.row_reference_setter, delimiter='\n    ')
    cursor_getter = self.foreach_col(
        ColumnSerializer.cursor_getter, delimiter='\n')
    cursor_setter = self.foreach_col(
        ColumnSerializer.cursor_setter, delimiter='\n')
    insert_argvalue = self.foreach_col(
        ColumnSerializer.row_insert_arg_value, delimiter=', ')
    return f'''
class {self.table_name} {{
 public:
  static constexpr auto kSpec = dataframe::CreateTypedDataframeSpec(
    {{{col_names}}},
    {col_specs});

  using Id = {self.table_name}_Id;
  struct RowReference;
  struct ConstRowReference;
  struct RowNumber {{
   public:
    explicit constexpr RowNumber(uint32_t value) : value_(value) {{}}
    uint32_t row_number() const {{ return value_; }}

    RowReference ToRowReference({self.table_name}* table) const {{
      return RowReference(table, value_);
    }}
    ConstRowReference ToRowReference(const {self.table_name}& table) const {{
      return ConstRowReference(&table, value_);
    }}

    bool operator==(const RowNumber& other) const {{
      return value_ == other.value_;
    }}
    bool operator<(const RowNumber& other) const {{
      return value_ < other.value_;
    }}
   private:
    uint32_t value_;
  }};
  struct ColumnIndex {{
    {col_index};
  }};
  struct RowReference {{
   public:
    explicit RowReference({self.table_name}* table, uint32_t row)
        : table_(table), row_(row) {{
        base::ignore_result(table_);
    }}
    {row_ref_getter}
    {row_ref_setter}
    RowNumber ToRowNumber() const {{
      return RowNumber{{row_}};
    }}

   private:
    friend struct ConstRowReference;
    {self.table_name}* table_;
    uint32_t row_;
  }};
  struct ConstRowReference {{
   public:
    explicit ConstRowReference(const {self.table_name}* table, uint32_t row)
        : table_(table), row_(row) {{
        base::ignore_result(table_);
    }}
    ConstRowReference(const RowReference& other)
        : table_(other.table_), row_(other.row_) {{}}
    {row_ref_getter}
    RowNumber ToRowNumber() const {{
      return RowNumber{{row_}};
    }}
   private:
    const {self.table_name}* table_;
    uint32_t row_;
  }};
  class ConstCursor {{
   public:
    explicit ConstCursor(const dataframe::Dataframe& df,
                         std::vector<dataframe::FilterSpec> filters,
                         std::vector<dataframe::SortSpec> sorts)
      : dataframe_(&df), cursor_(&df, std::move(filters), std::move(sorts)) {{
      base::ignore_result(dataframe_);
    }}

    PERFETTO_ALWAYS_INLINE void Execute() {{ cursor_.ExecuteUnchecked(); }}
    PERFETTO_ALWAYS_INLINE bool Eof() const {{ return cursor_.Eof(); }}
    PERFETTO_ALWAYS_INLINE void Next() {{ cursor_.Next(); }}
    template <typename C>
    PERFETTO_ALWAYS_INLINE void SetFilterValueUnchecked(uint32_t index, C value) {{
      cursor_.SetFilterValueUnchecked(index, std::move(value));
    }}
    PERFETTO_ALWAYS_INLINE void SetFilterValueListUnchecked(
        uint32_t index,
        const dataframe::TypedCursor::FilterValue* values,
        uint32_t count) {{
      cursor_.SetFilterValueListUnchecked(index, values, count);
    }}
    RowNumber ToRowNumber() const {{
      return RowNumber{{cursor_.RowIndex()}};
    }}
    void Reset() {{ cursor_.Reset(); }}
    {cursor_getter}

   private:
    const dataframe::Dataframe* dataframe_;
    dataframe::TypedCursor cursor_;
  }};
  class Cursor {{
   public:
    explicit Cursor(dataframe::Dataframe& df,
                    std::vector<dataframe::FilterSpec> filters,
                    std::vector<dataframe::SortSpec> sorts)
      : dataframe_(&df), cursor_(&df, std::move(filters), std::move(sorts)) {{
      base::ignore_result(dataframe_);
    }}

    PERFETTO_ALWAYS_INLINE void Execute() {{ cursor_.ExecuteUnchecked(); }}
    PERFETTO_ALWAYS_INLINE bool Eof() const {{ return cursor_.Eof(); }}
    PERFETTO_ALWAYS_INLINE void Next() {{ cursor_.Next(); }}
    template <typename C>
    PERFETTO_ALWAYS_INLINE void SetFilterValueUnchecked(uint32_t index, C value) {{
      cursor_.SetFilterValueUnchecked(index, std::move(value));
    }}
    PERFETTO_ALWAYS_INLINE void SetFilterValueListUnchecked(
        uint32_t index,
        const dataframe::TypedCursor::FilterValue* values,
        uint32_t count) {{
      cursor_.SetFilterValueListUnchecked(index, values, count);
    }}
    RowNumber ToRowNumber() const {{
      return RowNumber{{cursor_.RowIndex()}};
    }}
    void Reset() {{ cursor_.Reset(); }}

    {cursor_getter}
    {cursor_setter}

   private:
    dataframe::Dataframe* dataframe_;
    dataframe::TypedCursor cursor_;
  }};
  class Iterator {{
    public:
      explicit Iterator({self.table_name}* table) : table_(table) {{
        base::ignore_result(table_);
      }}
      explicit operator bool() const {{ return row_ < table_->row_count(); }}
      Iterator& operator++() {{
        ++row_;
        return *this;
      }}
      RowNumber row_number() const {{
        return RowNumber{{row_}};
      }}
      RowReference ToRowReference() const {{
        return RowReference(table_, row_);
      }}
      {row_ref_getter}
      {row_ref_setter}

    private:
      {self.table_name}* table_;
      uint32_t row_ = 0;
  }};
  class ConstIterator {{
    public:
      explicit ConstIterator(const {self.table_name}* table) : table_(table) {{
        base::ignore_result(table_);
      }}
      explicit operator bool() const {{ return row_ < table_->row_count(); }}
      ConstIterator& operator++() {{
        ++row_;
        return *this;
      }}
      RowNumber row_number() const {{
        return RowNumber{{row_}};
      }}
      ConstRowReference ToRowReference() const {{
        return ConstRowReference(table_, row_);
      }}
      {row_ref_getter}

    private:
      const {self.table_name}* table_;
      uint32_t row_ = 0;
  }};
  struct IdAndRow {{
    Id id;
    RowNumber row_number;
    uint32_t row;
    RowReference row_reference;
  }};
  {self.row_struct()}

  explicit {self.table_name}(StringPool* pool)
      : dataframe_(
        dataframe::Dataframe::CreateFromTypedSpec(kSpec, pool)) {{}}

  template <typename = void>
  IdAndRow Insert(const Row& row) {{
    uint32_t row_count = dataframe_.row_count();
    dataframe_.InsertUnchecked(kSpec, {insert_argvalue});
    return IdAndRow{{Id{{row_count}}, RowNumber{{row_count}}, row_count, RowReference(this, row_count)}};
  }}

  uint32_t row_count() const {{
    return dataframe_.row_count();
  }}

  uint64_t mutations() const {{
    return dataframe_.mutations();
  }}

  ConstRowReference operator[](Id id) const {{
    return ConstRowReference(this, id.value);
  }}
  ConstRowReference operator[](uint32_t row) const {{
    return ConstRowReference(this, row);
  }}

  RowReference operator[](Id id) {{
    return RowReference(this, id.value);
  }}
  RowReference operator[](uint32_t row) {{
    return RowReference(this, row);
  }}

  // Returns an `Id` for `value` if it refers to a valid row in this table, or
  // std::nullopt otherwise. Use this whenever turning an untrusted integer
  // (e.g. from SQL or a decoded proto) into a typed `Id`.
  std::optional<Id> TryCastId(int64_t value) const {{
    if (value < 0 || value >= row_count()) {{
      return std::nullopt;
    }}
    return Id{{static_cast<uint32_t>(value)}};
  }}

  ConstCursor CreateCursor(
      std::vector<dataframe::FilterSpec> filters = {{}},
      std::vector<dataframe::SortSpec> sorts = {{}}) const {{
    return ConstCursor(dataframe_, std::move(filters), std::move(sorts));
  }}
  Cursor CreateCursor(
      std::vector<dataframe::FilterSpec> filters = {{}},
      std::vector<dataframe::SortSpec> sorts = {{}}) {{
    return Cursor(dataframe_, std::move(filters), std::move(sorts));
  }}

  Iterator IterateRows() {{ return Iterator(this); }}
  ConstIterator IterateRows() const {{ return ConstIterator(this); }}

  void Finalize() {{ dataframe_.Finalize(); }}

  void Clear() {{ dataframe_.Clear(); }}

  static const char* Name() {{
    return "{self.table.sql_name}";
  }}

  dataframe::Dataframe& dataframe() {{
    return dataframe_;
  }}
  const dataframe::Dataframe& dataframe() const {{
    return dataframe_;
  }}

 private:
  dataframe::Dataframe dataframe_;
}};
static_assert(sizeof({self.table_name}) == sizeof(dataframe::Dataframe));
'''


def serialize_header(ifdef_guard: str, tables: List[ParsedTable],
                     include_paths: List[str], fwd_header_path: str) -> str:
  """Serializes a table header file containing the given set of tables."""
  # Replace the backslash with forward slash when building on Windows.
  # Caused b/327985369 without the replace.
  include_paths_str = '\n'.join([f'#include "{i}"' for i in include_paths
                                ]).replace("\\", "/")
  tables_str = '\n\n'.join([TableSerializer(t).serialize() for t in tables])
  fwd_header_str = fwd_header_path.replace("\\", "/")

  return f'''\
#ifndef {ifdef_guard}
#define {ifdef_guard}

#include <cstdint>
#include <optional>
#include <tuple>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>

#include "perfetto/base/compiler.h"
#include "perfetto/public/compiler.h"
#include "src/trace_processor/containers/string_pool.h"
#include "src/trace_processor/tables/macros_internal.h"
#include "src/trace_processor/core/dataframe/dataframe.h"
#include "src/trace_processor/core/dataframe/specs.h"
#include "src/trace_processor/core/dataframe/typed_cursor.h"
#include "{fwd_header_str}"

{include_paths_str}

namespace perfetto::trace_processor::tables {{

{tables_str}

}}  // namespace perfetto::trace_processor::tables

#endif  // {ifdef_guard}
'''


def serialize_fwd_header(ifdef_guard: str, tables: List[ParsedTable],
                         include_paths: List[str]) -> str:
  """Serializes a forward declaration header with Id types only."""
  # Replace the backslash with forward slash when building on Windows.
  include_paths_str = '\n'.join([
      f'#include "{i}"  // IWYU pragma: export' for i in include_paths
  ]).replace("\\", "/")
  id_types = '\n\n'.join(
      [TableSerializer(t).serialize_id_type() for t in tables])
  fwd_decls = '\n'.join([f'class {t.table.class_name};' for t in tables])

  # Generate tuple of table classes for this header
  table_names = ', '.join([t.table.class_name for t in tables])
  tuple_name = ifdef_guard.replace('_H_', '').replace(
      'SRC_TRACE_PROCESSOR_TABLES_', '') + '_TABLES'

  return f'''\
#ifndef {ifdef_guard}
#define {ifdef_guard}

#include <cstdint>
#include <tuple>

#include "src/trace_processor/tables/macros_internal.h"  // IWYU pragma: export

{include_paths_str}

namespace perfetto::trace_processor::tables {{

{id_types}

{fwd_decls}

// Tuple of all table types in this header
using {tuple_name} = std::tuple<{table_names}>;

}}  // namespace perfetto::trace_processor::tables

#endif  // {ifdef_guard}
'''