{#
 # Copyright (C) 2019  Google LLC
 #
 # 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.
 #}
{# response handling macros #}

{% macro sample_header(sample, calling_form) %}
# Generated code. DO NOT EDIT!
#
# Snippet for {{ sample.rpc }}
# NOTE: This snippet has been automatically generated for illustrative purposes only.
# It may require modifications to work in your environment.

# To install the latest published package dependency, execute the following:
#   python3 -m pip install {{ sample.package_name }}
{% endmacro %}

{% macro print_string_formatting(string_list) %}

{% if string_list|length == 1 %}
"{{ string_list[0] | render_format_string }}"
{% elif string_list|length == 2 and string_list[0] == "%s" and string_list[1] == "$resp" %}
response
{% else %}
{# Note: This is the equivalent of render_format_string(string_list[0], string_list[1:] )
# See https://jinja.palletsprojects.com/en/3.0.x/api/#custom-filters #}
f"{{ string_list[0] | render_format_string(string_list[1:]) }}"
{% endif %}
{% endmacro %}

{% macro print_input_params(full_request) %}
{% with input_parameters = [] %}
  {% for request in full_request.request_list %}
    {% if request.body %}
      {% for element in request.body if element.input_parameter %}
        {% do input_parameters.append(element.input_parameter) %}
      {% endfor %}
    {% elif request.single and request.single.input_parameter %}
      {% do input_parameters.append(request.single.input_parameter) %}
    {% endif %}
  {% endfor %}
{{ input_parameters|join(", ") -}}
{% endwith %}
{% endmacro %}

{# First elment is a format string, remaining elements are the format string parameters #}
{# Validating that the number of format params equals #}
{# the number of remaining params is handled by real python code #}
{% macro render_print(elts) %}
print({{ print_string_formatting(elts)|trim }})
{% endmacro %}

{% macro render_comment(elts) %}
  {# First elment is a format string, remaining elements are the format string parameters #}
  {# Validating that the number of format params equals #}
  {# the number of remaining params is handled by real python code #}
  {% with fmtStr = elts[0] %}
  {% with params = elts[1:]|map("coerce_response_name")|list %}
# {{ fmtStr|format(*params) }}
{% endwith %}
{% endwith %}
{% endmacro %}

{% macro render_define(statement) %}
{# Python code already verified the form, no need to check #}
{% with lvalue, rvalue = statement.split("=") %}
{{ lvalue }} = {{ rvalue|coerce_response_name }}
{% endwith %}
{% endmacro %}

{% macro render_collection_loop(statement) %}
for {{ statement.variable }} in {{ statement.collection|coerce_response_name }}:
  {% for s in statement.body %}
{{ dispatch_statement(s, 4) }}
  {% endfor %}
{% endmacro %}

{% macro render_map_loop(statement) %}
 {# At least one of key and value exist; validated in python #}
{% if "key" not in statement %}
for {{ statement.value }} in {{ statement.map|coerce_response_name }}.values():
{% elif "value" not in statement %}
for {{ statement.key }} in {{ statement.map|coerce_response_name }}.keys():
{% else %}
for {{statement.key }}, {{ statement.value }} in {{ statement.map|coerce_response_name }}.items():
{% endif %}
{% for s in statement.body %}
{{ dispatch_statement(s, 4) }}
{% endfor %}
{% endmacro %}

{% macro render_write_file(statement) %}
  {% with contents_rval = statement["contents"]|coerce_response_name %}
with open({{ print_string_formatting(statement["filename"])|trim }}, "wb") as f:
    f.write({{ contents_rval }})
  {% endwith %}
{% endmacro %}

{% macro dispatch_statement(statement, indentation=0) %}
{# Each statement is a dict with a single key/value pair #}
{% if "print" in statement %}
{{ render_print(statement["print"])|indent(width=indentation, first=True) }}
{% elif "define" in statement %}
{{ render_define(statement["define"])|indent(width=indentation, first=True) }}
{% elif "comment" in statement %}
{{ render_comment(statement["comment"])|indent(width=indentation, first=True) }}
{% elif "loop" in statement %}
  {% with loop = statement["loop"] %}
    {% if "collection" in loop %}
{{ render_collection_loop(loop)|indent(width=indentation, first=True) }}
    {% else %}
{{ render_map_loop(loop)|indent(width=indentation, first=True) }}
    {% endif %}
  {% endwith %}
{% elif "write_file" in statement %}
{{ render_write_file(statement["write_file"])|indent(indentation, first=True) }}
{% endif %}
{% endmacro %}

{% macro render_request_attr(base_name, attr) %}
{# Note: python code will have manipulated the value #}
{# to be the correct enum from the right module, if necessary. #}
{# Python is also responsible for verifying that each input parameter is unique, #}
{# no parameter is a reserved keyword #}
{% if attr.input_parameter %}

# {{ attr.input_parameter }} = {{ attr.value }}
{% if attr.value_is_file %}
with open({{ attr.input_parameter }}, "rb") as f:
    {{ base_name }}.{{ attr.field }} = f.read()
{% else %}
{{ base_name }}.{{ attr.field }} = {{ attr.input_parameter }}
{% endif %}
{% else %}
{{ base_name }}.{{ attr.field }} = {{ attr.value }}
{% endif %}
{% endmacro %}


{% macro render_client_setup(module_name, client_name) %}
# Create a client
client = {{ module_name }}.{{ client_name }}()
{% endmacro %}

{% macro render_request_setup(full_request, module_name, request_type, calling_form, calling_form_enum) %}
# Initialize request argument(s)
{% for parameter_block in full_request.request_list if parameter_block.body %}
{% if parameter_block.pattern %}
{# This is a resource-name patterned lookup parameter #}
{% with formals = [] %}
{% for attr in parameter_block.body %}
{{ attr.field }} = {{ attr.input_parameter or attr.value }}
{% endfor %}
{{ parameter_block.base }} = f"{{parameter_block.pattern }}"

{% endwith %}
{% else %}{# End resource name construction #}
{{ parameter_block.base }} = {{ module_name }}.{{ request_type.get_field(parameter_block.base).type.name }}()
{% for attr in parameter_block.body %}
{{ render_request_attr(parameter_block.base, attr) -}}
{% endfor %}

{% endif %}
{% endfor %}
{% if not full_request.flattenable %}
request = {{ module_name }}.{{ request_type.ident.name }}(
{% for parameter in full_request.request_list %}
    {{ parameter.base }}={{ parameter.base if parameter.body else parameter.single.value }},
{% endfor %}
)
{# Note: This template assumes only one request needs to be sent. When samples accept
configs the client streaming logic should be modified to allow 2+ request objects. #}
{# If client streaming, wrap the single request in a generator that produces 'requests' #}
{% if calling_form in [calling_form_enum.RequestStreamingBidi,
                       calling_form_enum.RequestStreamingClient] %}

# This method expects an iterator which contains
# '{{module_name}}.{{ request_type.ident.name }}' objects
# Here we create a generator that yields a single `request` for 
# demonstrative purposes.
requests = [request]

def request_generator():
    for request in requests:
        yield request
{% endif %}
{% endif %}
{% endmacro %}

{% macro render_request_params(request) %}
  {# Provide the top level parameters last and as keyword params #}
  {% with params = [] %}
    {% for r in request if r.body %}
      {% do params.append(r.base) %}
    {% endfor %}
    {% for r in request if r.single %}
      {% do params.append("%s=%s"|format(r.base, r.single.value)) %}
    {% endfor %}
{{ params|join(", ") -}}
  {% endwith %}
{% endmacro %}

{% macro render_request_params_unary(request) %}
  {# Provide the top level parameters last and as keyword params #}
  {% if request.flattenable %}
  {% with params = [] %}
    {% for r in request.request_list %}
      {% do params.append("%s=%s"|format(r.base, r.single.value if r.single else r.base)) %}
    {% endfor %}
{{ params|join(", ") -}}
  {% endwith %}
  {% else %}
request=request
  {% endif %}
{% endmacro %}


{% macro render_method_call(sample, calling_form, calling_form_enum, transport) %}
{# Note: this doesn't deal with enums or unions #}
{# LROs return operation objects and paged requests return pager objects #}
{% if transport == "grpc-async" and calling_form not in
[calling_form_enum.LongRunningRequestPromise, calling_form_enum.RequestPagedAll] %}
await{{ " "}}
{%- endif -%}
{% if calling_form in [calling_form_enum.RequestStreamingBidi,
                       calling_form_enum.RequestStreamingClient] %}
client.{{ render_method_name(sample)|trim }}(requests=request_generator())
{% else %}{# TODO: deal with flattening #}
{# TODO: set up client streaming once some questions are answered #}
client.{{ render_method_name(sample)|trim }}({{ render_request_params_unary(sample.request)|trim }})
{% endif %}
{% endmacro %}

{% macro operation_text(transport) %}
{% if transport == "grpc-async" %}
(await operation)
{% else %}
operation
{% endif %}
{% endmacro %}

{# Setting up the method invocation is the responsibility of the caller: #}
{# it's just easier to set up client side streaming and other things from outside this macro. #}
{% macro render_calling_form(method_invocation_text, calling_form, calling_form_enum, transport, response_statements ) %}
# Make the request
{% if calling_form in [calling_form_enum.Request, calling_form_enum.RequestStreamingClient] %}
{% if response_statements %}response = {% endif %}{{ method_invocation_text|trim }}

{% if response_statements %}
# Handle the response
{% for statement in response_statements %}
{{ dispatch_statement(statement)|trim }}
{% endfor %}
{% endif %}
{% elif calling_form == calling_form_enum.RequestPagedAll %}
page_result = {{ method_invocation_text|trim }}

# Handle the response
{% if transport == "grpc-async" %}async {% endif %}for response in page_result:
  {% for statement in response_statements %}
    {{ dispatch_statement(statement)|trim }}
  {% endfor %}
{% elif calling_form == calling_form_enum.RequestPaged %}
page_result = {{ method_invocation_text|trim }}

# Handle the response
{% if transport == "grpc-async" %}async {% endif %}for page in page_result.pages():
    for response in page:
  {% for statement in response_statements %}
        {{ dispatch_statement(statement)|trim }}
  {% endfor %}
{% elif calling_form in [calling_form_enum.RequestStreamingServer,
                        calling_form_enum.RequestStreamingBidi] %}
stream = {{ method_invocation_text|trim }}

# Handle the response
{% if transport == "grpc-async" %}async {% endif %}for response in stream:
  {% for statement in response_statements %}
    {{ dispatch_statement(statement)|trim }}
  {% endfor %}
{% elif calling_form == calling_form_enum.LongRunningRequestPromise %}
operation = {{ method_invocation_text|trim }}

print("Waiting for operation to complete...")

{% with operation_text = operation_text(transport) %}
response = {{ operation_text|trim }}.result()
{% endwith %}

# Handle the response
{% for statement in response_statements %}
{{ dispatch_statement(statement)|trim }}
{% endfor %}
{% endif %}
{% endmacro %}

{% macro render_method_name(sample) %}
{% if sample.is_internal %}
_{{ sample.rpc|snake_case }}
{% else %}
{{ sample.rpc|snake_case }}
{% endif %}
{% endmacro %}

{% macro render_main_block(sample, full_request) %}
def main():
    import argparse

    parser = argparse.ArgumentParser()
{% with arg_list = [] %}
{% for request in full_request.request_list if request.body %}
{% for attr in request.body if attr.input_parameter %}
    parser.add_argument("--{{ attr.input_parameter }}",
                        type=str,
                        default={{ attr.value }})
{% do arg_list.append("args." + attr.input_parameter) %}
{% endfor %}
{% endfor %}
{% for request in full_request.request_list if request.single and request.single.input_parameter %}
    parser.add_argument("--{{ request.single.input_parameter }}",
                        type=str,
                        default={{ request.single.value }})
{% endfor %}
    args = parser.parse_args()

    sample_{{ render_method_name(sample)|trim }}({{ arg_list|join(", ") }})


if __name__ == "__main__":
    main()
{% endwith %}
{% endmacro %}
