""" pytest for WorkflowTasks
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# pylint: disable=protected-access,missing-docstring,invalid-name

from mock import MagicMock
import pytest

from DIRAC import gLogger
from DIRAC.Interfaces.API.Job import Job

# sut
from DIRAC.TransformationSystem.Client.WorkflowTasks import WorkflowTasks

gLogger.setLevel("DEBUG")

mockTransClient = MagicMock()
mockTransClient.setTaskStatusAndWmsID.return_value = {"OK": True}

WMSClientMock = MagicMock()
jobMonitoringClient = MagicMock()

wfTasks = WorkflowTasks(
    transClient=mockTransClient,
    submissionClient=WMSClientMock,
    jobMonitoringClient=jobMonitoringClient,
    outputDataModule="mock",
)
odm_o = MagicMock()
odm_o.execute.return_value = {"OK": True, "Value": {}}
wfTasks.outputDataModule_o = odm_o

taskDict = {
    1: {"TransformationID": 1, "a1": "aa1", "b1": "bb1", "Site": "MySite"},
    2: {"TransformationID": 1, "a2": "aa2", "b2": "bb2", "InputData": ["a1", "a2"]},
    3: {"TransformationID": 2, "a3": "aa3", "b3": "bb3"},
}

taskDictSimple = {1: {"TransformationID": 1, "Site": "MySite"}, 2: {"TransformationID": 1}, 3: {"TransformationID": 2}}

taskDictSimpleInputs = {1: {"TransformationID": 1, "InputData": ["a1", "a2", "a3"]}}

taskDictNoInputs = {
    1: {"TransformationID": 1, "a1": "aa1", "b1": "bb1", "Site": "MySite"},
    2: {"TransformationID": 1, "a2": "aa2", "b2": "bb2"},
    3: {"TransformationID": 2, "a3": "aa3", "b3": "bb3"},
}

taskDictNoInputsNoSite = {1: {"TransformationID": 1}, 2: {"TransformationID": 1}, 3: {"TransformationID": 2}}

expected = {
    "OK": True,
    "Value": {
        1: {"a1": "aa1", "TaskObject": "", "TransformationID": 1, "b1": "bb1", "Site": "MySite", "JobType": "User"},
        2: {
            "TaskObject": "",
            "a2": "aa2",
            "TransformationID": 1,
            "InputData": ["a1", "a2"],
            "b2": "bb2",
            "Site": "MySite",
            "JobType": "User",
        },
        3: {"TaskObject": "", "a3": "aa3", "TransformationID": 2, "b3": "bb3", "Site": "MySite", "JobType": "User"},
    },
}

expectedBulkSimple = {
    "OK": True,
    "Value": {"BulkJobObject": "", 1: {"TransformationID": 1, "InputData": ["a1", "a2", "a3"], "JobType": "User"}},
}

expectedBulk = {
    "OK": True,
    "Value": {
        "BulkJobObject": "",
        1: {"a1": "aa1", "TransformationID": 1, "b1": "bb1", "Site": "MySite", "JobType": "User"},
        2: {"a2": "aa2", "TransformationID": 1, "b2": "bb2", "InputData": ["a1", "a2"], "JobType": "User"},
        3: {"TransformationID": 2, "a3": "aa3", "b3": "bb3", "JobType": "User"},
    },
}


@pytest.mark.parametrize(
    "taskDictionary, bulkSubmissionFlag, result, expectedRes",
    [
        (taskDict, False, True, expected),
        (taskDict, True, False, expectedBulk),
        (taskDictSimple, True, True, expectedBulk),
        (taskDictSimpleInputs, True, True, expectedBulkSimple),
        (taskDictNoInputs, True, False, expectedBulk),
        (taskDictNoInputsNoSite, True, True, expectedBulk),
    ],
)
def test_prepareTranformationTasks(taskDictionary, bulkSubmissionFlag, result, expectedRes):
    res = wfTasks.prepareTransformationTasks(
        "", taskDictionary, "test_user", "test_group", "test_DN", bulkSubmissionFlag=bulkSubmissionFlag
    )
    assert res["OK"] == result
    if res["OK"]:
        for key, value in res["Value"].items():
            if key != "BulkJobObject":
                assert key in expectedRes["Value"]
                for tKey, tValue in value.items():
                    assert tKey in expectedRes["Value"][key]
                    if tKey == "TaskObject" and tValue:
                        assert isinstance(tValue, Job)
                    else:
                        assert tValue == expectedRes["Value"][key][tKey]


def ourgetSitesForSE(ses):
    if ses == ["pippo"] or ses == "pippo":
        return {"OK": True, "Value": ["Site1"]}
    elif ses == ["pluto"] or ses == "pluto":
        return {"OK": True, "Value": ["Site2"]}
    elif ses == ["pippo", "pluto"] or ses == "pippo,pluto":
        return {"OK": True, "Value": ["Site1", "Site2"]}


@pytest.mark.parametrize(
    "paramsDict, expected",
    [
        ({"Site": "", "TargetSE": ""}, ["ANY"]),
        ({"Site": "ANY", "TargetSE": ""}, ["ANY"]),
        ({"TargetSE": "Unknown"}, ["ANY"]),
        ({"Site": "Site2", "TargetSE": ""}, ["Site2"]),
        ({"Site": "Site1;Site2", "TargetSE": "pippo"}, ["Site1"]),
        ({"Site": "Site1;Site2", "TargetSE": "pippo,pluto"}, ["Site1", "Site2"]),
        ({"Site": "Site1;Site2;Site3", "TargetSE": "pippo,pluto"}, ["Site1", "Site2"]),
        ({"Site": "Site2", "TargetSE": "pippo,pluto"}, ["Site2"]),
        ({"Site": "ANY", "TargetSE": "pippo,pluto"}, ["Site1", "Site2"]),
        ({"Site": "Site1", "TargetSE": "pluto"}, []),
    ],
)
def test__handleDestination(mocker, paramsDict, expected):
    mocker.patch("DIRAC.TransformationSystem.Client.TaskManagerPlugin.getSitesForSE", side_effect=ourgetSitesForSE)
    res = wfTasks._handleDestination(paramsDict)
    assert sorted(res) == sorted(expected)
