#######################################################################################
# Copyright ETSI Contributors and Others.
#
# 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 unittest.mock import Mock

import asynctest
from mock import AsyncMock, patch
from osm_common.dataclasses.temporal_dataclasses import (
    GetVnfRecordIdsInput,
    GetVnfRecordIdsOutput,
    LcmOperationState,
    ModelInfo,
    NsLcmOperationInput,
    NsState,
    UpdateLcmOperationStateInput,
    UpdateNsStateInput,
    VnfInstantiateInput,
)
from osm_common.temporal_constants import (
    ACTIVITY_CREATE_MODEL,
    ACTIVITY_GET_VNF_RECORD_IDS,
    ACTIVITY_UPDATE_LCM_OPERATION_STATE,
    ACTIVITY_UPDATE_NS_STATE,
    LCM_TASK_QUEUE,
    WORKFLOW_VNF_INSTANTIATE,
)
from temporalio import activity, workflow
from temporalio.client import WorkflowFailureError
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

from osm_lcm.temporal.ns_workflows import NsInstantiateWorkflow

vnfr_ids = [
    "828d91ee-fa04-43bb-8471-f66ea74597e7",
    "c4bbeb41-df7e-4daa-863d-c8fd29fac96d",
]


class TestException(Exception):
    pass


@activity.defn(name=ACTIVITY_CREATE_MODEL)
async def mock_create_model(create_model_input: ModelInfo) -> None:
    pass


@activity.defn(name=ACTIVITY_GET_VNF_RECORD_IDS)
async def mock_get_vnf_record_ids(
    get_vnf_record_ids_input: GetVnfRecordIdsInput,
) -> None:
    return GetVnfRecordIdsOutput(vnfr_ids=vnfr_ids)


@activity.defn(name=ACTIVITY_CREATE_MODEL)
async def mock_create_model_raises(create_model_input: ModelInfo) -> None:
    raise TestException("Test exception")


@workflow.defn(name=WORKFLOW_VNF_INSTANTIATE, sandboxed=False)
class MockVnfInstantiateWorkflow:
    @workflow.run
    async def run(self, input: VnfInstantiateInput) -> None:
        pass


class TestNsInstantiateWorkflow(asynctest.TestCase):
    input = NsLcmOperationInput(
        nslcmop={
            "_id": "1234",
            "nsInstanceId": "5678",
            "operationParams": {"vimAccountId": "9876"},
        }
    )

    @activity.defn(name=ACTIVITY_UPDATE_LCM_OPERATION_STATE)
    async def mock_update_lcm_operation_state(
        self,
        data: UpdateLcmOperationStateInput,
    ) -> None:
        self.mock_update_lcm_operation_state_tracker(data)

    @activity.defn(name=ACTIVITY_UPDATE_NS_STATE)
    async def mock_update_ns_state(self, data: UpdateNsStateInput) -> None:
        self.mock_update_ns_state_tracker(data)

    async def setUp(self):
        self.env = await WorkflowEnvironment.start_time_skipping()
        self.mock_update_lcm_operation_state_tracker = Mock()
        self.mock_update_ns_state_tracker = Mock()

    @patch("temporalio.workflow.execute_child_workflow")
    async def test_instantiate_workflow(self, mock_execute_child_workflow: AsyncMock):
        async with self.env as env:
            async with Worker(
                env.client,
                task_queue=LCM_TASK_QUEUE,
                workflows=[NsInstantiateWorkflow, MockVnfInstantiateWorkflow],
                activities=[
                    mock_create_model,
                    mock_get_vnf_record_ids,
                    self.mock_update_ns_state,
                    self.mock_update_lcm_operation_state,
                ],
                debug_mode=True,
            ):
                await env.client.execute_workflow(
                    NsInstantiateWorkflow,
                    arg=self.input,
                    id=self.input.nslcmop["nsInstanceId"],
                    task_queue=LCM_TASK_QUEUE,
                )

        assert [
            call.kwargs["workflow"] == WORKFLOW_VNF_INSTANTIATE
            and call.kwargs["arg"].vnfr_uuid
            for call in mock_execute_child_workflow.call_args_list
        ] == vnfr_ids

        assert_lcm_op_states(
            self.mock_update_lcm_operation_state_tracker.call_args_list,
            [LcmOperationState.PROCESSING, LcmOperationState.COMPLETED],
        )

        assert_ns_states(
            self.mock_update_ns_state_tracker.call_args_list, [NsState.INSTANTIATED]
        )

    async def test_instantiate_workflow_with_exception_updates_lcm_operation(self):
        async with self.env as env:
            async with Worker(
                env.client,
                task_queue=LCM_TASK_QUEUE,
                workflows=[NsInstantiateWorkflow, MockVnfInstantiateWorkflow],
                activities=[
                    mock_create_model_raises,
                    mock_get_vnf_record_ids,
                    self.mock_update_ns_state,
                    self.mock_update_lcm_operation_state,
                ],
                debug_mode=True,
            ):
                # TODO: Check this WorkflowFailureError is caused by a TestException
                with self.assertRaises(WorkflowFailureError):
                    await env.client.execute_workflow(
                        NsInstantiateWorkflow,
                        arg=self.input,
                        id=self.input.nslcmop["nsInstanceId"],
                        task_queue=LCM_TASK_QUEUE,
                    )

        assert_lcm_op_states(
            self.mock_update_lcm_operation_state_tracker.call_args_list,
            [LcmOperationState.PROCESSING, LcmOperationState.FAILED],
        )

    async def test_instantiate_workflow_with_exception_updates_ns_state(self):
        async with self.env as env:
            async with Worker(
                env.client,
                task_queue=LCM_TASK_QUEUE,
                workflows=[NsInstantiateWorkflow, MockVnfInstantiateWorkflow],
                activities=[
                    mock_create_model_raises,
                    mock_get_vnf_record_ids,
                    self.mock_update_ns_state,
                    self.mock_update_lcm_operation_state,
                ],
                debug_mode=True,
            ):
                # TODO: Check this WorkflowFailureError is caused by a TestException
                with self.assertRaises(WorkflowFailureError):
                    await env.client.execute_workflow(
                        NsInstantiateWorkflow,
                        arg=self.input,
                        id=self.input.nslcmop["nsInstanceId"],
                        task_queue=LCM_TASK_QUEUE,
                    )
        assert_ns_states(
            self.mock_update_ns_state_tracker.call_args_list, [NsState.INSTANTIATED]
        )


def assert_ns_states(call_args_list, expected_states):
    """Asserts that the NS state was set to each of the expected states (in order)."""
    assert [call.args[0].state for call in call_args_list] == expected_states


def assert_lcm_op_states(call_args_list, expected_states):
    """Asserts that the LCM operation was set to each of the exepcted states (in order)."""
    assert [call.args[0].op_state for call in call_args_list] == expected_states
