from flows_sdk.blocks import PythonBlock
from flows_sdk.flows import Flow, Manifest, Parameter
from flows_sdk.implementations.idp_v42.additional_blocks import (
    HyperscienceRestApiBlock,
    JSONOutputsBlock,
)
from flows_sdk.package_utils import export_flow
from flows_sdk.types import HsTask
from flows_sdk.utils import str_to_deterministic_uuid_4, workflow_input

SUBMISSION_AWARE_ON_ERROR_FLOW_IDENTIFIER = 'SUBMISSION_AWARE_ON_ERROR_FLOW_V2_EXAMPLE'
SUBMISSION_AWARE_ON_ERROR_FLOW_UUID = str_to_deterministic_uuid_4(
    SUBMISSION_AWARE_ON_ERROR_FLOW_IDENTIFIER
)
SUBMISSION_AWARE_ON_ERROR_FLOW_TITLE = 'On-Error with included Submission data V2 flows-sdk example'
SUBMISSION_AWARE_ON_ERROR_DESCRIPTION = """
Flow designed to be used as an "on-error" handler for submission-based flows.
\n
With the on-error flow configured, when the original flow fails for any reason, and all
    the configured auto-retries have been exhausted, the on-error flow will be triggered.
\n
Version V2 of the flow uses an endpoint for fetching the flow run that is much lighter and
    faster - it returns the metadata and status of the flow run, but omits heavy block input/output
    payloads that could burden the system further in case of systematic failure. Old versions of
    the on-error flow use a now-deprecated API endpoint, and will stop working in V45 of the
    product, unless updated to use the new lighter endpoint - /api/v5/flow_runs/id/summary.
\n
In this on-error flow implementation,
    we fetch the failed flow run, alongside with the full submission metadata
    (if the submission ID was found), and then send it to a user-configured output.
    The data format of the output conforms to
    the schema documented here: https://docs.hyperscience.ai/#flows-runs and
    here: https://docs.hyperscience.ai/#submissions, under the respective keys
    "flow_run" and "submission".
\n
If you need more complex logic in this flow
    please consult the documentation of Flows SDK (https://flows-sdk.hyperscience.ai)
    and the available public APIs (https://docs.hyperscience.ai)
""".lstrip()


class OnErrorFlowInputKeys:
    failed_run_uuid: str = 'failed_run_uuid'


class OnErrorManifest(Manifest):
    def __init__(self, flow_identifier: str):
        super().__init__(
            identifier=flow_identifier,
            roles=['on_error', 'supporting'],
            input=[
                Parameter(
                    name=OnErrorFlowInputKeys.failed_run_uuid,
                    type='string',
                    title='Failed Flow Run UUID',
                    description='UUID of the failed flow run that triggered this on_error flow',
                    optional=False,
                    ui={'hidden': True},
                )
            ],
        )


def submission_aware_on_error_flow() -> Flow:
    flow_run_rest_api_block = HyperscienceRestApiBlock(
        reference_name='flow_run_rest_api_block',
        title='Get Flow Run',
        description='Get the data of the failed flow run',
        method='GET',
        app_endpoint=f'/api/v5/flow_runs/{workflow_input(OnErrorFlowInputKeys.failed_run_uuid)}/summary',
    )

    def _get_correlation_id(_hs_task: HsTask) -> dict:
        return {'correlation_id': _hs_task.correlation_id}

    get_correlation_id_block = PythonBlock(
        reference_name='get_correlation_id',
        title='Get Correlation ID',
        description='Get the correlation ID for the failed Submission / Flow Run',
        code=_get_correlation_id,
    )

    fetch_submissions_block = HyperscienceRestApiBlock(
        reference_name='fetch_submission',
        title='Fetch Submission',
        description='Fetch the submission data for the failed flow run',
        method='GET',
        app_endpoint=f'/api/v5/submissions'
        '?debug=true'
        f'&correlation_id={get_correlation_id_block.output("correlation_id")}',
    )

    def _prep_output_data(flow_run: dict, get_submissions_response: dict) -> dict:
        submissions_in_response = get_submissions_response.get('results', [])
        submission = None if not submissions_in_response else submissions_in_response[0]
        return {'flow_run': flow_run, 'submission': submission}

    prep_output_data_block = PythonBlock(
        title='Format Submission Response',
        description='Prepare the output data to be sent to the user-configured output'
        'under the keys "flow_run" and "submission"',
        code=_prep_output_data,
        code_input={
            'flow_run': flow_run_rest_api_block.output('result.data'),
            'get_submissions_response': fetch_submissions_block.output('result.data'),
        },
    )

    json_outputs_block = JSONOutputsBlock(inputs={'payload': prep_output_data_block.output()})

    return Flow(
        uuid=SUBMISSION_AWARE_ON_ERROR_FLOW_UUID,
        owner_email='flows-sdk@hyperscience.ai',
        title=SUBMISSION_AWARE_ON_ERROR_FLOW_TITLE,
        description=SUBMISSION_AWARE_ON_ERROR_DESCRIPTION,
        manifest=OnErrorManifest(SUBMISSION_AWARE_ON_ERROR_FLOW_IDENTIFIER),
        output={},
        input={},
        blocks=[
            flow_run_rest_api_block,
            get_correlation_id_block,
            fetch_submissions_block,
            prep_output_data_block,
            json_outputs_block,
        ],
    )


if __name__ == '__main__':
    export_flow(submission_aware_on_error_flow())
