"""Test the Transformationinfo."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import OrderedDict
from contextlib import contextmanager

import pytest
from mock import MagicMock as Mock, patch

from DIRAC import S_OK, S_ERROR
import DIRAC
import DIRAC.TransformationSystem.Client.TransformationClient
import DIRAC.Resources.Catalog.FileCatalogClient
import DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient
from DIRAC.TransformationSystem.Utilities.JobInfo import JobInfo
from DIRAC.TransformationSystem.Utilities.TransformationInfo import TransformationInfo

from DIRAC.tests.Utilities.utils import MatchStringWith

__RCSID__ = "$Id$"

# pylint: disable=W0212, redefined-outer-name


@pytest.fixture
def userProxyFixture(mocker):
  """Mock UserProxy."""
  @contextmanager
  def _mockedCM(*args, **kwargs):
    try:
      yield S_OK()
    finally:
      pass
  mocker.patch('DIRAC.TransformationSystem.Utilities.TransformationInfo.UserProxy', new=_mockedCM)


@pytest.fixture
def failingUserProxyFixture(mocker):
  """Mock UserProxy."""
  @contextmanager
  def _mockedCM(*args, **kwargs):
    try:
      yield S_ERROR("Failed to set up proxy")
    finally:
      pass
  mocker.patch('DIRAC.TransformationSystem.Utilities.TransformationInfo.UserProxy', new=_mockedCM)


@pytest.fixture
def tiFixture():
  """Fixture for TransformationInfo object."""
  tMock = Mock(name="transMock", spec=DIRAC.TransformationSystem.Client.TransformationClient.TransformationClient)
  tMock.setFileStatusForTransformation = Mock(name="setFileStat")
  fcMock = Mock(name="fcMock", spec=DIRAC.Resources.Catalog.FileCatalogClient.FileCatalogClient)
  jmMock = Mock(name="jobMonMock", spec=DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient.JobMonitoringClient)
  jsucMock = Mock(name='jsuc', spec=DIRAC.WorkloadManagementSystem.Client.JobStateUpdateClient.JobStateUpdateClient)
  transInfoDict = dict(TransformationID=1234, TransformationName="TestProd12", Type="TestProd",
                       AuthorDN='/some/cert/owner', AuthorGroup='Test_Prod')

  tri = TransformationInfo(transformationID=1234, transInfoDict=transInfoDict,
                           enabled=False,
                           tClient=tMock, fcClient=fcMock, jobMon=jmMock)
  tri.log = Mock(name="LogMock")
  tri.jobStateClient = jsucMock
  return tri


@pytest.fixture
def tdFixture():
  """Fixture for tasksDicts."""
  taskDicts = [dict(TaskID=123, LFN="lfn123", Status="Assigned", FileID=987001, ErrorCount=9),
               dict(TaskID=124, LFN="lfn124", Status="Processed", FileID=987002, ErrorCount=8),
               ]
  return taskDicts


def test_init(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo init..........................."""
  assert isinstance(tiFixture, TransformationInfo)
  assert not tiFixture.enabled


def test_checkTasksStatus(tiFixture, tdFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo checkTasksStatus..............."""
  # error getting files
  tiFixture.tClient.getTransformationFiles.return_value = S_ERROR("nope")
  with pytest.raises(RuntimeError) as re:
    tiFixture.checkTasksStatus()
  assert "Failed to get transformation tasks: nope" in str(re)

  # success getting files
  tiFixture.tClient.getTransformationFiles.return_value = S_OK(tdFixture)
  retDict = tiFixture.checkTasksStatus()
  assert len(retDict) == 2
  assert 123 in retDict
  assert 124 in retDict
  assert len(retDict[124]) == 1
  assert 'FileID' in retDict[124][0]


def test_setJob_Status(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo setJob functions..............."""
  job = Mock(spec=JobInfo)
  job.jobID = 5678
  tiFixture.enabled = False
  tiFixture._TransformationInfo__setTaskStatus = Mock()
  tiFixture._TransformationInfo__updateJobStatus = Mock()

  # setJobDone
  tiFixture.setJobDone(job)
  tiFixture._TransformationInfo__setTaskStatus.assert_not_called()
  tiFixture._TransformationInfo__updateJobStatus.assert_not_called()

  tiFixture.enabled = True
  job.status = "Done"
  tiFixture.setJobDone(job)
  tiFixture._TransformationInfo__setTaskStatus.assert_called_once_with(job, "Done")
  tiFixture._TransformationInfo__updateJobStatus.assert_not_called()

  tiFixture.enabled = True
  job.status = "Failed"
  tiFixture._TransformationInfo__setTaskStatus.reset_mock()
  tiFixture.setJobDone(job)
  tiFixture._TransformationInfo__setTaskStatus.assert_called_once_with(job, "Done")
  tiFixture._TransformationInfo__updateJobStatus.assert_called_once_with(job.jobID, 'Done', "Job forced to Done")

  # setJobFailed
  tiFixture.enabled = False
  tiFixture._TransformationInfo__setTaskStatus.reset_mock()
  tiFixture._TransformationInfo__updateJobStatus.reset_mock()
  tiFixture.setJobFailed(job)
  tiFixture._TransformationInfo__setTaskStatus.assert_not_called()
  tiFixture._TransformationInfo__updateJobStatus.assert_not_called()

  tiFixture.enabled = True
  job.status = "Failed"
  tiFixture._TransformationInfo__setTaskStatus.reset_mock()
  tiFixture._TransformationInfo__updateJobStatus.reset_mock()
  tiFixture.setJobFailed(job)
  tiFixture._TransformationInfo__setTaskStatus.assert_called_once_with(job, "Failed")
  tiFixture._TransformationInfo__updateJobStatus.assert_not_called()

  tiFixture.enabled = True
  job.status = "Done"
  tiFixture._TransformationInfo__setTaskStatus.reset_mock()
  tiFixture._TransformationInfo__updateJobStatus.reset_mock()
  tiFixture.setJobFailed(job)
  tiFixture._TransformationInfo__setTaskStatus.assert_called_once_with(job, "Failed")
  tiFixture._TransformationInfo__updateJobStatus.assert_called_once_with(job.jobID, "Failed", "Job forced to Failed")


def test_setInputStatusFuncs(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo setInput functions............."""
  tiFixture._TransformationInfo__setInputStatus = Mock()
  tiFixture.enabled = False
  job = Mock(spec=JobInfo)

  tiFixture._TransformationInfo__setInputStatus.reset_mock()
  tiFixture.setInputUnused(job)
  tiFixture._TransformationInfo__setInputStatus.assert_called_once_with(job, "Unused")

  tiFixture._TransformationInfo__setInputStatus.reset_mock()
  tiFixture.setInputProcessed(job)
  tiFixture._TransformationInfo__setInputStatus.assert_called_once_with(job, "Processed")

  tiFixture._TransformationInfo__setInputStatus.reset_mock()
  tiFixture.setInputDeleted(job)
  tiFixture._TransformationInfo__setInputStatus.assert_called_once_with(job, "Deleted")

  tiFixture._TransformationInfo__setInputStatus.reset_mock()
  tiFixture.setInputMaxReset(job)
  tiFixture._TransformationInfo__setInputStatus.assert_called_once_with(job, "MaxReset")


def test_setInputStatus(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo setInputStatus................."""
  job = Mock(spec=JobInfo)
  job.inputFiles = ['dummylfn', 'otherDummy']
  status = "Unused"

  tiFixture.enabled = False
  tiFixture._TransformationInfo__setInputStatus(job, status)

  tiFixture.enabled = True
  tiFixture.tClient.setFileStatusForTransformation.return_value = S_ERROR("Failed to set")
  with pytest.raises(RuntimeError) as re:
    tiFixture._TransformationInfo__setInputStatus(job, status)
  assert "Failed updating file status" in str(re)

  tiFixture.enabled = True
  tiFixture.tClient.setFileStatusForTransformation.return_value = S_OK("All Good")
  tiFixture._TransformationInfo__setInputStatus(job, status)


def test_setTaskStatus(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo setTaskStatus.................."""
  job = Mock(spec=JobInfo)
  job.taskID = 1234
  tiFixture.tClient.setTaskStatus = Mock(return_value=S_OK("Done"))
  tiFixture._TransformationInfo__setTaskStatus(job, "Processed")
  tiFixture.tClient.setTaskStatus.assert_called_once_with(tiFixture.transName, 1234, "Processed")

  tiFixture.tClient.setTaskStatus = Mock(return_value=S_ERROR("NotDone"))
  with pytest.raises(RuntimeError) as re:
    tiFixture._TransformationInfo__setTaskStatus(job, "Processed")
  tiFixture.tClient.setTaskStatus.assert_called_once_with(tiFixture.transName, 1234, "Processed")
  assert "Failed updating task status: NotDone" in str(re)


def test_updateJobStatus(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo updateJobStatus................"""

  tiFixture.jobStateClient.setJobStatus = Mock()
  tiFixture.jobStateClient.setJobStatus.return_value = S_OK()
  tiFixture.enabled = False
  res = tiFixture._TransformationInfo__updateJobStatus(1234, 'Failed', minorstatus=None)
  assert res['OK']
  assert res['Value'] == 'DisabledMode'
  tiFixture.jobStateClient.setJobStatus.assert_not_called()

  tiFixture.jobStateClient.setJobStatus.reset()
  tiFixture.jobStateClient.setJobStatus.return_value = S_OK('added record')
  tiFixture.enabled = True
  res = tiFixture._TransformationInfo__updateJobStatus(1234, 'Failed', minorstatus=None)
  assert res['OK']
  assert res['Value'] == 'added record'
  tiFixture.jobStateClient.setJobStatus.assert_called_once_with(1234, 'Failed', None, 'DataRecoveryAgent')

  tiFixture.jobStateClient.setJobStatus.return_value = S_ERROR('Error setting job status')
  tiFixture.enabled = True
  with pytest.raises(RuntimeError) as re:
    tiFixture._TransformationInfo__updateJobStatus(1234, 'Failed', minorstatus=None)
  assert 'Failed to update job status' in str(re)


def test_findAllDescendants(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo findAllDescendents............."""
  tiFixture.fcClient.getFileDescendents = Mock(return_value=S_OK({"Successful": {"lfn1": ["lfnD1", "lfnD2"],
                                                                                 "lfnD1": ["lfnDD1", "lfnDD2"],
                                                                                 }}))
  descList = tiFixture._TransformationInfo__findAllDescendants(lfnList=[])
  assert set(descList) == {"lfnDD1", "lfnDD2", "lfnD1", "lfnD2"}

  tiFixture.fcClient.getFileDescendents = Mock(return_value=S_ERROR("Cannot get descendants"))
  descList = tiFixture._TransformationInfo__findAllDescendants(lfnList=[])
  assert descList == []


def test_cleanOutputs_proxyFail(tiFixture, failingUserProxyFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo cleanOutputs..................."""
  descList = ["lfnDD1", "lfnDD2", "lfnD1", "lfnD2"]

  jobInfo = Mock(spec=JobInfo)
  jobInfo.outputFiles = ["lfn1", "lfn2"]
  jobInfo.outputFileStatus = ["Exists", "Missing"]

  tiFixture.enabled = True
  tiFixture._TransformationInfo__findAllDescendants = Mock(return_value=descList)
  with pytest.raises(RuntimeError) as re:
    tiFixture.cleanOutputs(jobInfo)
  assert "Failed to get a proxy" in str(re)


def test_cleanOutputs(tiFixture, userProxyFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo cleanOutputs..................."""
  descList = ["lfnDD1", "lfnDD2", "lfnD1", "lfnD2"]

  jobInfo = Mock(spec=JobInfo)
  jobInfo.outputFiles = ["lfn1", "lfn2"]
  jobInfo.outputFileStatus = ["Exists", "Missing"]

  tiFixture.enabled = False
  tiFixture._TransformationInfo__findAllDescendants = Mock(return_value=descList)

  tiFixture.cleanOutputs(jobInfo)
  tiFixture.log.notice.assert_any_call(MatchStringWith("Would have removed these files"))
  tiFixture.log.notice.assert_any_call(MatchStringWith("lfn1"))
  for _name, args, _kwargs in tiFixture.log.notice.mock_calls:
    assert "lfn2" not in str(args)

  remMock = Mock(name="remmock")
  remMock.removeFile.return_value = S_ERROR("arg")

  tiFixture.enabled = True
  tiFixture._TransformationInfo__findAllDescendants = Mock(return_value=descList)
  with patch("DIRAC.TransformationSystem.Utilities.TransformationInfo.DataManager",
             return_value=remMock,
             autospec=True):
    with pytest.raises(RuntimeError) as re:
      tiFixture.cleanOutputs(jobInfo)
    assert "Failed to remove LFNs: arg" in str(re)

  remMock = Mock(name="remmock")
  remMock.removeFile.return_value = S_OK({"Successful": {"lfn1": "OK", "lfn2": "OK"},
                                          "Failed": {"lfnD2": "SomeReason",
                                                     "lfnD3": "SomeReason",
                                                     "lfnDD2": "SomeOtherReason"}})

  tiFixture.enabled = True
  with patch("DIRAC.TransformationSystem.Utilities.TransformationInfo.DataManager",
             autospec=True,
             return_value=remMock):
    tiFixture.cleanOutputs(jobInfo)
    tiFixture.log.notice.assert_any_call(MatchStringWith("Successfully removed 2 files"))

  # nothing to remove
  jobInfo = Mock(spec=JobInfo)
  jobInfo.outputFiles = []
  tiFixture._TransformationInfo__findAllDescendants = Mock(return_value=descList)
  tiFixture.cleanOutputs(jobInfo)
  tiFixture._TransformationInfo__findAllDescendants.assert_not_called()

  # nothing to remove
  jobInfo = Mock(spec=JobInfo)
  jobInfo.outputFiles = ["lfn1", "lfn2"]
  jobInfo.outputFileStatus = ["Missing", "Missing"]
  tiFixture._TransformationInfo__findAllDescendants = Mock(return_value=[])
  tiFixture.cleanOutputs(jobInfo)
  tiFixture._TransformationInfo__findAllDescendants.assert_called_once_with(jobInfo.outputFiles)


def test_getJobs(tiFixture):
  """DIRAC.TransformationSystem.Utilities.TransformationInfo getJobs........................"""
  tiFixture.jobMon.getJobs = Mock()
  tiFixture.jobMon.getJobs.side_effect = (S_OK([123, 456, 789]), S_OK([1123, 1456, 1789]))

  # All OK, just Done
  jobs, ndone, nfailed = tiFixture.getJobs(statusList=["Done"])
  attrDict = dict(Status=["Done"], JobGroup="00001234")
  tiFixture.jobMon.getJobs.assert_called_once_with(attrDict)
  assert ndone == 3
  assert nfailed == 0
  assert isinstance(jobs, OrderedDict)

  # All OK, just Done
  tiFixture.jobMon.getJobs.reset_mock()
  tiFixture.jobMon.getJobs.side_effect = (S_OK([123, 456, 789]), S_OK([1123, 1456, 1789]))
  jobs, ndone, nfailed = tiFixture.getJobs(statusList=["Failed"])
  attrDict = dict(Status=["Failed"], JobGroup="00001234")
  tiFixture.jobMon.getJobs.assert_called_once_with(attrDict)
  assert ndone == 0
  assert nfailed == 3
  assert isinstance(jobs, OrderedDict)

  # All OK, None
  tiFixture.jobMon.getJobs.reset_mock()
  tiFixture.jobMon.getJobs.side_effect = (S_OK([123, 456, 789]), S_OK([1123, 56, 89]))
  jobs, ndone, nfailed = tiFixture.getJobs(statusList=None)
  attrDict = dict(Status=["Done"], JobGroup="00001234")
  attrDict2 = dict(Status=["Failed"], JobGroup="00001234")
  tiFixture.jobMon.getJobs.assert_any_call(attrDict)
  tiFixture.jobMon.getJobs.assert_any_call(attrDict2)
  assert ndone == 3
  assert nfailed == 3
  assert isinstance(jobs, OrderedDict)
  assert [56, 89, 123, 456, 789, 1123] == list(jobs)

  # All ERROR
  tiFixture.jobMon.getJobs = Mock()
  tiFixture.jobMon.getJobs.side_effect = (S_ERROR("Not Done"), S_ERROR("Not Failed"))
  with pytest.raises(RuntimeError) as re:
    jobs, ndone, nfailed = tiFixture.getJobs(statusList=None)
  attrDict = dict(Status=["Done"], JobGroup="00001234")
  tiFixture.jobMon.getJobs.assert_called_once_with(attrDict)
  assert "Failed to get jobs" in str(re)
