# Copyright 2021 Agnostiq Inc.
#
# This file is part of Covalent.
#
# Licensed under the Apache License 2.0 (the "License"). A copy of the
# License may be obtained with this software package or at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Use of this file is prohibited except in compliance with the License.
# 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.

"""Unit tests for AWS ECS executor."""

import os
import shutil
from pathlib import Path
from unittest import mock
from unittest.mock import AsyncMock

import cloudpickle as pickle
import pytest

from covalent_ecs_plugin.ecs import FUNC_FILENAME, RESULT_FILENAME, ECSExecutor


class TestECSExecutor:
    MOCK_PROFILE = "my_profile"
    MOCK_S3_BUCKET_NAME = "s3-bucket"
    MOCK_ECS_CLUSTER_NAME = "ecs-cluster"
    MOCK_ECS_EXECUTION_ROLE = "task-execution-role"
    MOCK_ECS_TASK_ROLE_NAME = "task-role-name"
    MOCK_ECS_TASK_SUBNET_ID = "sb-1234"
    MOCK_ECS_TASK_SG_ID = "sg-123"
    MOCK_ECS_LOG_GROUP_NAME = "log-group"
    MOCK_VCPU = 0.1234
    MOCK_MEMORY = "123"
    MOCK_POLL_FREQ = 123
    MOCK_DISPATCH_ID = 112233
    MOCK_NODE_ID = 1
    MOCK_TASK_ARN = "task-arn/123"

    @property
    def MOCK_FUNC_FILENAME(self):
        return FUNC_FILENAME.format(dispatch_id=self.MOCK_DISPATCH_ID, node_id=self.MOCK_NODE_ID)

    @property
    def MOCK_RESULT_FILENAME(self):
        return RESULT_FILENAME.format(dispatch_id=self.MOCK_DISPATCH_ID, node_id=self.MOCK_NODE_ID)

    @property
    def MOCK_TASK_METADATA(self):
        return {"dispatch_id": self.MOCK_DISPATCH_ID, "node_id": self.MOCK_NODE_ID}

    @pytest.fixture
    def mock_executor_config(self, tmp_path):
        MOCK_CREDENTIALS_FILE: Path = tmp_path / "credentials"
        MOCK_CREDENTIALS_FILE.touch()
        return {
            "profile": self.MOCK_PROFILE,
            "s3_bucket_name": self.MOCK_S3_BUCKET_NAME,
            "ecs_cluster_name": self.MOCK_ECS_CLUSTER_NAME,
            "ecs_task_execution_role_name": self.MOCK_ECS_EXECUTION_ROLE,
            "ecs_task_role_name": self.MOCK_ECS_TASK_ROLE_NAME,
            "ecs_task_subnet_id": self.MOCK_ECS_TASK_SUBNET_ID,
            "ecs_task_security_group_id": self.MOCK_ECS_TASK_SG_ID,
            "ecs_task_log_group_name": self.MOCK_ECS_LOG_GROUP_NAME,
            "vcpu": self.MOCK_VCPU,
            "memory": self.MOCK_MEMORY,
            "poll_freq": self.MOCK_POLL_FREQ,
        }

    @pytest.fixture
    def mock_executor(self, mock_executor_config):
        executor = ECSExecutor(**mock_executor_config)
        executor.cache_dir = "/tmp/covalent_test_cache_dir/"
        return executor

    @mock.patch.dict(os.environ, {}, clear=True)
    def test_init_explicit_values(self, mocker, mock_executor_config):
        """Test executor class values are overridden properly during instantiation"""

        # only call to get_config is get_config("executors.ecs.cache_dir")
        mocker.patch("covalent_ecs_plugin.ecs.get_config", return_value="mock_cache_dir")
        mocker.patch("covalent_ecs_plugin.ecs.ECSExecutor._is_valid_subnet_id", return_value=True)
        mocker.patch(
            "covalent_ecs_plugin.ecs.ECSExecutor._is_valid_security_group", return_value=True
        )

        executor = ECSExecutor(**mock_executor_config)

        assert executor.profile == self.MOCK_PROFILE
        assert executor.s3_bucket_name == self.MOCK_S3_BUCKET_NAME
        assert executor.execution_role == self.MOCK_ECS_EXECUTION_ROLE
        assert executor.ecs_task_role_name == self.MOCK_ECS_TASK_ROLE_NAME
        assert executor.ecs_task_subnet_id == self.MOCK_ECS_TASK_SUBNET_ID
        assert executor.ecs_task_security_group_id == self.MOCK_ECS_TASK_SG_ID
        assert executor.log_group_name == self.MOCK_ECS_LOG_GROUP_NAME
        assert executor.vcpu == self.MOCK_VCPU
        assert executor.memory == self.MOCK_MEMORY
        assert executor.poll_freq == self.MOCK_POLL_FREQ

    @pytest.mark.asyncio
    async def test_upload_file_to_s3(self, mock_executor, mocker):
        """Test to upload file to s3."""
        Path(mock_executor.cache_dir).mkdir(parents=True, exist_ok=True)

        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")

        def some_function():
            pass

        mock_executor._upload_task_to_s3(
            some_function,
            self.MOCK_DISPATCH_ID,
            self.MOCK_NODE_ID,
            ("some_arg"),
            {"some": "kwarg"},
        )
        boto3_mock.Session().client().upload_file.assert_called_once()

        shutil.rmtree(mock_executor.cache_dir)

    @pytest.mark.asyncio
    async def test_upload_task(self, mock_executor, mocker):
        """Test for method to call the upload task method."""

        def some_function(x):
            return x

        upload_to_s3_mock = mocker.patch(
            "covalent_ecs_plugin.ecs.ECSExecutor._upload_task_to_s3", return_value=AsyncMock()
        )

        await mock_executor._upload_task(some_function, (1), {}, self.MOCK_TASK_METADATA)
        upload_to_s3_mock.assert_called_once_with(
            self.MOCK_DISPATCH_ID, self.MOCK_NODE_ID, some_function, (1), {}
        )

    @pytest.mark.asyncio
    async def test_submit_task(self, mock_executor, mocker):
        """Test submit task method."""
        MOCK_IDENTITY = {"Account": 1234}
        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")
        await mock_executor.submit_task(self.MOCK_TASK_METADATA, MOCK_IDENTITY)
        boto3_mock.Session().client().register_task_definition.assert_called_once()
        boto3_mock.Session().client().run_task.assert_called_once()

    def test_is_valid_subnet_id(self, mock_executor):
        """Test the valid subnet checking method."""
        assert mock_executor._is_valid_subnet_id("subnet-871545e1") is True
        assert mock_executor._is_valid_subnet_id("subnet-871545e") is False
        assert mock_executor._is_valid_subnet_id("jlkjlkj871545e1") is False

    def test_is_valid_security_group(self, mock_executor):
        """Test the valid security group checking method."""
        assert mock_executor._is_valid_security_group("sg-0043541a") is True
        assert mock_executor._is_valid_security_group("sg-0043541") is False
        assert mock_executor._is_valid_security_group("80980043541") is False

    @pytest.mark.asyncio
    async def test_cancel(self, mocker, mock_executor):
        """Test the execution cancellation method."""
        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")
        await mock_executor.cancel("mock_task_arn", "mock_reason")
        boto3_mock.Session().client().stop_task.assert_called_once_with(
            cluster=self.MOCK_ECS_CLUSTER_NAME, task="mock_task_arn", reason="mock_reason"
        )

    @pytest.mark.asyncio
    async def test_query_result(self, mocker, mock_executor, tmp_path: Path):
        """Test the method to query the result."""
        Path(mock_executor.cache_dir).mkdir(parents=True, exist_ok=True)
        mock_local_result_path = Path(mock_executor.cache_dir) / self.MOCK_RESULT_FILENAME
        mock_local_result_path.touch()

        MOCK_RESULT_CONTENTS = "mock_result"

        with open(mock_local_result_path, "wb") as f:
            pickle.dump(MOCK_RESULT_CONTENTS, f)

        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")
        result = await mock_executor.query_result(task_metadata=self.MOCK_TASK_METADATA)
        assert result == MOCK_RESULT_CONTENTS

        boto3_mock.Session().client().download_file.assert_called_once_with(
            self.MOCK_S3_BUCKET_NAME, self.MOCK_RESULT_FILENAME, str(mock_local_result_path)
        )
        shutil.rmtree(mock_executor.cache_dir)

    @pytest.mark.asyncio
    async def test_get_log_events(self, mocker, mock_executor):
        """Test the method to retrieve the log events from the log stream."""

        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")
        boto3_mock.Session().client().get_log_events.return_value = {
            "events": [{"message": "hello"}, {"message": "world"}]
        }
        log_events = await mock_executor._get_log_events(
            "task-arn", task_metadata=self.MOCK_TASK_METADATA
        )
        assert log_events == "hello\nworld\n"

    @pytest.mark.asyncio
    async def test_get_status(self, mocker, mock_executor):
        """Test the status checking method."""
        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")
        ecs_client_mock = boto3_mock.Session().client()

        # Case 1: no tasks found
        ecs_client_mock.get_paginator().paginate.return_value = []
        res = await mock_executor.get_status(self.MOCK_TASK_ARN)
        assert res == ("TASK_NOT_FOUND", -1)

        # Case 2 valid task found
        ecs_client_mock.get_paginator().paginate.return_value = [
            {"taskArns": [self.MOCK_TASK_ARN]}
        ]
        ecs_client_mock.describe_tasks.return_value = {
            "tasks": [
                {
                    "taskArn": self.MOCK_TASK_ARN,
                    "lastStatus": "RUNNING",
                    "containers": [{"exitCode": 1}],
                }
            ]
        }
        res = await mock_executor.get_status(self.MOCK_TASK_ARN)
        assert res == ("RUNNING", 1)

        # Case 3 - task found without any status
        ecs_client_mock.get_paginator().paginate.return_value = [
            {"taskArns": [self.MOCK_TASK_ARN]}
        ]
        ecs_client_mock.describe_tasks.return_value = {
            "tasks": [{"taskArn": self.MOCK_TASK_ARN, "lastStatus": "FAILED"}]
        }
        res = await mock_executor.get_status(self.MOCK_TASK_ARN)
        assert res == ("FAILED", -1)

    @pytest.mark.asyncio
    async def test_poll_ecs_task(self, mocker, mock_executor):
        """Test the method to poll the ecs task."""

        mock_executor.poll_freq = self.MOCK_POLL_FREQ
        sleep_mock = mocker.patch("covalent_ecs_plugin.ecs.asyncio.sleep")

        mocker.patch(
            "covalent_ecs_plugin.ecs.ECSExecutor.get_status",
            side_effect=[("RUNNING", 1), ("STOPPED", 0)],
        )
        await mock_executor._poll_task(self.MOCK_TASK_ARN)
        sleep_mock.assert_called_once_with(self.MOCK_POLL_FREQ)

        with pytest.raises(Exception):
            await mock_executor._poll_task(self.MOCK_TASK_ARN)

    @pytest.mark.asyncio
    async def test_run(self, mocker, mock_executor):
        """Test the run method."""

        MOCK_IDENTITY = {"Account": 1234}
        mock_executor.vcpu = 1
        mock_executor.memory = 1

        def mock_func(x):
            return x

        boto3_mock = mocker.patch("covalent_ecs_plugin.ecs.boto3")

        upload_task_mock = mocker.patch("covalent_ecs_plugin.ecs.ECSExecutor._upload_task")
        validate_credentials_mock = mocker.patch(
            "covalent_ecs_plugin.ecs.ECSExecutor._validate_credentials"
        )
        submit_task_mock = mocker.patch("covalent_ecs_plugin.ecs.ECSExecutor.submit_task")
        _poll_task_mock = mocker.patch("covalent_ecs_plugin.ecs.ECSExecutor._poll_task")
        query_result_mock = mocker.patch("covalent_ecs_plugin.ecs.ECSExecutor.query_result")

        validate_credentials_mock.return_value = MOCK_IDENTITY

        await mock_executor.run(
            function=mock_func, args=[], kwargs={"x": 1}, task_metadata=self.MOCK_TASK_METADATA
        )

        upload_task_mock.assert_called_once_with(mock_func, [], {"x": 1}, self.MOCK_TASK_METADATA)
        validate_credentials_mock.assert_called_once()
        submit_task_mock.assert_called_once_with(self.MOCK_TASK_METADATA, MOCK_IDENTITY)

        returned_task_arn = await submit_task_mock()

        _poll_task_mock.assert_called_once_with(returned_task_arn)
        query_result_mock.assert_called_once_with(self.MOCK_TASK_METADATA)
        shutil.rmtree(mock_executor.cache_dir)
