{# Modify workflow -#}

{% from 'macros.j2' import required with context -%}

from typing import Annotated

import structlog
from pydantic import AfterValidator, ConfigDict, model_validator
from pydantic_forms.types import FormGenerator, State, UUIDstr

{%- if use_updated_readonly_field is true  %}
from pydantic_forms.validators import read_only_field
{%- else %}
from pydantic_forms.validators import ReadOnlyField
{%- endif %}

from orchestrator.forms import FormPage
from orchestrator.forms.validators import CustomerId, Divider
from orchestrator.types import SubscriptionLifecycle
from orchestrator.workflow import StepList, begin, step
from orchestrator.workflows.steps import set_status
from orchestrator.workflows.utils import modify_workflow

{%- for namespace, type in types_to_import %}
from {{ namespace }} import {{ type }}
{%- endfor %}
from {{ product_types_module }}.{{ product.variable }} import {{ product.type }}, {{ product.type }}Provisioning
{% if validation_imports -%}
from {{workflows_module}}.{{ product.variable }}.shared.forms import {{ validation_imports | join(", ") }}
{%- endif %}
{% if product.config.summary_forms -%}
from {{workflows_module}}.shared import modify_summary_form
{%- endif %}
{% if product_block_types -%}
from {{ product_blocks_module }}.{{ product_block.variable }} import {{ product_block_types | map(attribute='type') | join(", ") }}
{%- endif %}

{% include 'additional_modify_imports.j2' %}


logger = structlog.get_logger(__name__)

{% for validation in validations -%}
validated_{{ validation.field.name }} = Annotated[{{ validation.field.type }}, AfterValidator({{ validation.validation.id }}_validator)]
{% endfor %}

def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
    subscription = {{ product.type }}.from_subscription(subscription_id)
    {{ product_block.name }} = subscription.{{ product_block.name }}

    # TODO fill in additional fields if needed

    class Modify{{ product.type }}Form(FormPage):
        customer_id: CustomerId = subscription.customer_id  # type: ignore
{% include "additional_modify_input_fields.j2" ignore missing %}

        divider_1: Divider

        {% for field in fields if not field.modifiable is defined -%}
        {% if use_updated_readonly_field is true  -%}
        {{ field.name }}: read_only_field({{ product_block.name }}.{{ field.name }})
        {% else -%}
        {{ field.name }}: ReadOnlyField({{ product_block.name }}.{{ field.name }})
        {% endif -%}
        {% endfor -%}

        {% for field in fields if field.modifiable is defined -%}
        {% if field.type in validations | map(attribute='field.type') -%}
            {% set type = 'validated_' +  field.name -%}
        {% else -%}
            {% set type = field.type -%}
        {% endif -%}
        {{ field.name }}: {{ required(field, type, ["inactive", "provisioning"]) }} = {{ product_block.name }}.{{ field.name }}
        {% endfor -%}
        {%- for validation in workflow_validations %}

        @model_validator(mode="after")
        def {{ validation.id }}(self) -> "Modify{{ product.type }}Form":
            if False:  # TODO implement validation
                raise ValueError("{{ validation.description }}")
            return self
        {%- endfor %}

    user_input = yield Modify{{ product.type }}Form
    user_input_dict = user_input.dict()
    {% if product.config.summary_forms %}
    summary_fields = [{% for field in fields %}{% if field.type not in existing_product_blocks %}"{{ field.name }}", {% endif %}{% endfor %}]
    yield from modify_summary_form(user_input_dict, subscription.{{ product_block.name }}, summary_fields)
    {%- endif %}

    return user_input_dict | {"subscription": subscription}


@step("Update subscription")
def update_subscription(
    subscription: {{ product.type }}Provisioning,
    {% for field in fields if field.modifiable is defined -%}
    {% if field.type not in existing_product_blocks -%}
    {{ field.name }}: {{ required(field, field.type, ["inactive", "provisioning"]) }},
    {%- endif %}
    {% endfor -%}) -> State:
    # TODO: get all modified fields
    {% for field in fields if field.modifiable is defined -%}
    subscription.{{ product_block.name }}.{{ field.name }} = {{ field.name }}
    {% endfor %}
    return {"subscription": subscription}


@step("Update subscription description")
def update_subscription_description(subscription: {{ product.type }}) -> State:
    subscription.description = subscription_description(subscription)
    return {"subscription": subscription}


{% include 'additional_create_steps.j2' %}


@modify_workflow("Modify {{ product.name }}", initial_input_form=initial_input_form_generator, additional_steps=additional_steps)
def modify_{{ product.variable }}() -> StepList:
    return (
        begin
        >> set_status(SubscriptionLifecycle.PROVISIONING)
        >> update_subscription
        >> update_subscription_description
        # TODO add additional steps if needed
        >> set_status(SubscriptionLifecycle.ACTIVE)
    )
