# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.

import argparse
import doctest
import os
from subprocess import check_output
from typing import Dict, List, TextIO


THRIFT_HEADER = """
# This file is generated by `fbcode/thrift/test/testset:generator`
# {'@' + 'generated'}

namespace cpp2 apache.thrift.test.testset
"""

FIELD_COUNT = 2  # Number of fields per structs


def format_dict(
    d: Dict[str, str], key_format: str, value_format: str
) -> Dict[str, str]:
    """Format key/value of dict
    >>> result = format_dict({"foo_k": "foo_v", "bar_k": "bar_v"}, 'prefix_{}', "{}_suffix")
    >>> result == {'prefix_foo_k': 'foo_v_suffix', 'prefix_bar_k': 'bar_v_suffix'}
    True
    """
    return {key_format.format(k): value_format.format(d[k]) for k in d}


PRIMITIVE_TYPES = [
    "bool",
    "byte",
    "i16",
    "i32",
    "i64",
    "float",
    "double",
    "binary",
    "string",
]


def generate_union_names_to_types() -> Dict[str, str]:
    """ Generate display name to thrift type mapping in union. Display name will be used in file name, rule name, etc """
    ret = {t: t for t in PRIMITIVE_TYPES}
    ret.update(format_dict(ret, "set_{}", "set<{}>"))
    ret.update(format_dict(ret, "map_string_{}", "map<string, {}>"))
    ret.update(format_dict(ret, "{}_cpp_ref", "{} (cpp.ref = 'true')"))
    return ret


def generate_struct_names_to_types() -> Dict[str, str]:
    """ Similar to thrift types in union. Difference is that unions cannot contain qualified fields. """
    ret = generate_union_names_to_types()
    ret.update(
        **format_dict(ret, "optional_{}", "optional {}"),
        **format_dict(ret, "required_{}", "required {}"),
    )
    return ret


def generate_class(class_type: str, name: str, types: List[str]) -> str:
    """Generate thrift struct from types
    >>> print(generate_class("struct", "Foo", ["i64", "optional string", "set<i32> (cpp.ref = 'true')"]))
    struct Foo {
      1: i64 field_1;
      2: optional string field_2;
      3: set<i32> (cpp.ref = 'true') field_3;
    }
    """
    lines = [f"{class_type} {name} {{"]
    for i, t in enumerate(types):
        lines.append("  {0}: {1} field_{0};".format(i + 1, t))
    lines.append(f'}} (any_type.name="facebook.com/thrift/test/testset/{name}")')
    return "\n".join(lines)


def print_thrift_class(
    file: TextIO, class_type: str, names_to_types: Dict[str, str]
) -> None:
    name = "empty_" + class_type
    print(generate_class(class_type, name, []), file=file)
    classes = [name]
    for display_name, type in names_to_types.items():
        class_name = class_type + "_" + display_name
        classes.append(class_name)
        print(generate_class(class_type, class_name, [type] * FIELD_COUNT), file=file)

    # Thrift class that contains all other generated classes with same-type
    print(generate_class(class_type, class_type + "_all", classes), file=file)


def gen_struct_all(path: str) -> None:
    with open(path, "w") as file:
        print(THRIFT_HEADER, file=file)
        print_thrift_class(file, "struct", generate_struct_names_to_types())
        print_thrift_class(file, "union", generate_union_names_to_types())


def main() -> None:
    doctest.testmod()
    os.chdir(check_output(["buck", "root"]).strip())
    parser = argparse.ArgumentParser()
    parser.add_argument("--install_dir", required=True)
    parser.add_argument("--filename", required=True)
    args = parser.parse_args()
    gen_struct_all(os.path.join(args.install_dir, args.filename))


if __name__ == "__main__":
    main()

# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.

import argparse
import doctest
import os
from subprocess import check_output
from typing import Dict, List, TextIO


THRIFT_HEADER = """
# This file is generated by `fbcode/thrift/test/testset:generator`
# {'@' + 'generated'}

namespace cpp2 apache.thrift.test.testset
"""

FIELD_COUNT = 2  # Number of fields per structs


def format_dict(
    d: Dict[str, str], key_format: str, value_format: str
) -> Dict[str, str]:
    """Format key/value of dict
    >>> result = format_dict({"foo_k": "foo_v", "bar_k": "bar_v"}, 'prefix_{}', "{}_suffix")
    >>> result == {'prefix_foo_k': 'foo_v_suffix', 'prefix_bar_k': 'bar_v_suffix'}
    True
    """
    return {key_format.format(k): value_format.format(d[k]) for k in d}


PRIMITIVE_TYPES = [
    "bool",
    "byte",
    "i16",
    "i32",
    "i64",
    "float",
    "double",
    "binary",
    "string",
]


def generate_union_names_to_types() -> Dict[str, str]:
    """ Generate display name to thrift type mapping in union. Display name will be used in file name, rule name, etc """
    ret = {t: t for t in PRIMITIVE_TYPES}
    ret.update(format_dict(ret, "set_{}", "set<{}>"))
    ret.update(format_dict(ret, "map_string_{}", "map<string, {}>"))
    ret.update(format_dict(ret, "{}_cpp_ref", "{} (cpp.ref = 'true')"))
    return ret


def generate_struct_names_to_types() -> Dict[str, str]:
    """ Similar to thrift types in union. Difference is that unions cannot contain qualified fields. """
    ret = generate_union_names_to_types()
    ret.update(
        **format_dict(ret, "optional_{}", "optional {}"),
        **format_dict(ret, "required_{}", "required {}"),
    )
    return ret


def generate_class(class_type: str, name: str, types: List[str]) -> str:
    """Generate thrift struct from types
    >>> print(generate_class("struct", "Foo", ["i64", "optional string", "set<i32> (cpp.ref = 'true')"]))
    struct Foo {
      1: i64 field_1;
      2: optional string field_2;
      3: set<i32> (cpp.ref = 'true') field_3;
    }
    """
    lines = [f"{class_type} {name} {{"]
    for i, t in enumerate(types):
        lines.append("  {0}: {1} field_{0};".format(i + 1, t))
    lines.append(f'}} (any_type.name="facebook.com/thrift/test/testset/{name}")')
    return "\n".join(lines)


def print_thrift_class(
    file: TextIO, class_type: str, names_to_types: Dict[str, str]
) -> None:
    name = "empty_" + class_type
    print(generate_class(class_type, name, []), file=file)
    classes = [name]
    for display_name, type in names_to_types.items():
        class_name = class_type + "_" + display_name
        classes.append(class_name)
        print(generate_class(class_type, class_name, [type] * FIELD_COUNT), file=file)

    # Thrift class that contains all other generated classes with same-type
    print(generate_class(class_type, class_type + "_all", classes), file=file)


def gen_struct_all(path: str) -> None:
    with open(path, "w") as file:
        print(THRIFT_HEADER, file=file)
        print_thrift_class(file, "struct", generate_struct_names_to_types())
        print_thrift_class(file, "union", generate_union_names_to_types())


def main() -> None:
    doctest.testmod()
    os.chdir(check_output(["buck", "root"]).strip())
    parser = argparse.ArgumentParser()
    parser.add_argument("--install_dir", required=True)
    parser.add_argument("--filename", required=True)
    args = parser.parse_args()
    gen_struct_all(os.path.join(args.install_dir, args.filename))


if __name__ == "__main__":
    main()

# Copyright (c) Facebook, Inc. and its affiliates. # # 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. # This file is generated by `fbcode/thrift/test/testset:generator` # {'@' + 'generated'}  namespace cpp2 apache.thrift.test.testset # Number of fields per structs Format key/value of dict     >>> result = format_dict({"foo_k": "foo_v", "bar_k": "bar_v"}, 'prefix_{}', "{}_suffix")     >>> result == {'prefix_foo_k': 'foo_v_suffix', 'prefix_bar_k': 'bar_v_suffix'}     True Generate display name to thrift type mapping in union. Display name will be used in file name, rule name, etc Similar to thrift types in union. Difference is that unions cannot contain qualified fields. Generate thrift struct from types     >>> print(generate_class("struct", "Foo", ["i64", "optional string", "set<i32> (cpp.ref = 'true')"]))     struct Foo {       1: i64 field_1;       2: optional string field_2;       3: set<i32> (cpp.ref = 'true') field_3;     } # Thrift class that contains all other generated classes with same-type