from datetime import datetime
import json
import os
import tempfile
import unittest
from collections import Counter
from typing import List
from unittest.mock import MagicMock, Mock, patch

import b2luigi
from b2luigi.batch.processes.gbasf2 import Gbasf2Process, JobStatus, get_unique_project_name

from ..helpers import B2LuigiTestCase
from .batch_task_1 import MyTask


class MyGbasf2Task(MyTask):
    gbasf2_project_name_prefix = "my_gb2_task_"


class TestGbasf2RescheduleJobs(B2LuigiTestCase):
    job_statuses_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "_gbasf2_project_statuses")
    joblist_tmpfile_name = "jobs_to_be_rescheduled.txt"

    def setUp(self):
        super().setUp()
        self.gb2_mock_process = MagicMock()
        self.gb2_mock_process.task = MyGbasf2Task("some_parameter")
        self.gb2_mock_process.dirac_user = "username"
        self.gb2_mock_process.gbasf2_project_name = get_unique_project_name(self.gb2_mock_process.task)
        self.gb2_mock_process.n_retries_by_job = Counter()
        self.gb2_mock_process.max_retries = 1
        b2luigi.set_setting("gbasf2_print_status_updates", False)

    def _get_job_status_dict(self, job_status_fname):
        with open(os.path.join(self.job_statuses_dir, job_status_fname)) as job_status_json_file:
            return json.load(job_status_json_file)

    def _reschedule_jobs(self, job_ids):
        pass

    def assert_rescheduled_jobs(self, job_status_fname, expected_jobs_to_be_rescheduled):
        with patch(
            "b2luigi.batch.processes.gbasf2.get_gbasf2_project_job_status_dict",
            MagicMock(return_value=self._get_job_status_dict(job_status_fname)),
        ):
            with patch("b2luigi.batch.processes.gbasf2.Gbasf2Process._reschedule_jobs", new=self._reschedule_jobs):
                Gbasf2Process._reschedule_failed_jobs(self.gb2_mock_process)

                self.assertCountEqual(
                    list(self.gb2_mock_process.n_retries_by_job.keys()), expected_jobs_to_be_rescheduled
                )
                for jobid in self.gb2_mock_process.n_retries_by_job.keys():
                    self.assertEqual(self.gb2_mock_process.n_retries_by_job[jobid], 1)

    def test_reschedule_jobs_all_done(self):
        "Test gbasf2 project rescheduling with dict where all jobs are done"
        self.assert_rescheduled_jobs("done_testjbucket1357828d80b3.json", [])

    def test_reschedule_jobs_one_failed(self):
        "Test gbasf2 project rescheduling with dict where one job in project failed"
        self.assert_rescheduled_jobs("failed_7663_r03743_10_prod00013766_11x1.json", ["188623842", "188625261"])

    def test_reschedule_jobs_running(self):
        "Test gbasf2 project rescheduling with dict where several jobs are either running or waiting"
        self.assert_rescheduled_jobs("running_TrainingFEI_17-04-2021a10a957289.json", [])

    def test_reschedule_jobs_major_status_done_but_minor_status_not(self):
        "Test gbasf2 project rescheduling with dict where major job status is all done but application status is not ``DONE``"
        self.assert_rescheduled_jobs("all_done_but_application_error.json", ["187522107"])


class TestGbasf2GetJobStatus(B2LuigiTestCase):
    job_statuses_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "_gbasf2_project_statuses")

    def setUp(self):
        super().setUp()
        self.gb2_mock_process = Mock()
        self.gb2_mock_process.task = MyGbasf2Task("some_parameter")
        self.gb2_mock_process.dirac_user = "username"
        self.gb2_mock_process.gbasf2_project_name = get_unique_project_name(self.gb2_mock_process.task)
        self.gb2_mock_process.max_retries = 0
        self.gb2_mock_process._last_job_status_update = datetime.fromtimestamp(0)
        self.gb2_mock_process._job_status_max_age = 5 * 60
        b2luigi.set_setting("gbasf2_print_status_updates", False)

    def _get_job_status_dict(self, job_status_fname):
        with open(os.path.join(self.job_statuses_dir, job_status_fname)) as job_status_json_file:
            return json.load(job_status_json_file)

    def assert_job_status(self, job_status_fname, expected_job_status):
        with patch(
            "b2luigi.batch.processes.gbasf2.get_gbasf2_project_job_status_dict",
            MagicMock(return_value=self._get_job_status_dict(job_status_fname)),
        ):
            job_status = Gbasf2Process.get_job_status(self.gb2_mock_process)
            self.assertEqual(job_status, expected_job_status)

    def test_get_job_status_all_done(self):
        "Test gbasf2 project status dict where all jobs are done"
        self.assert_job_status("done_testjbucket1357828d80b3.json", JobStatus.successful)

    def test_get_job_status_one_failed(self):
        "Test gbasf2 project status dict where one job in project failed"
        self.assert_job_status("failed_7663_r03743_10_prod00013766_11x1.json", JobStatus.aborted)

    def test_get_job_status_running(self):
        "Test gbasf2 project status dict where several jobs are either running or waiting"
        self.assert_job_status("running_TrainingFEI_17-04-2021a10a957289.json", JobStatus.running)

    def test_get_job_status_major_status_done_but_minor_status_not(self):
        "Test gbasf2 project status dict where major job status is all done but application status is not ``DONE``"
        self.assert_job_status("all_done_but_application_error.json", JobStatus.aborted)


class TestBuildGbasf2SubmitCommand(B2LuigiTestCase):
    dummy_lfn = "/grid/path/to/dataset"

    def setUp(self):
        super().setUp()
        self.dummy_lfn_file = tempfile.NamedTemporaryFile()

    def _get_task_file_dir(self, task):
        return self.test_dir

    def _build_gbasf2_submit_command(self, task) -> List[str]:
        gb2_mock_process = Mock()
        os.makedirs(self._get_task_file_dir(task), exist_ok=True)
        self.pickle_file_path = os.path.join(self._get_task_file_dir(task), "serialized_basf2_path.pkl")
        self.wrapper_file_path = os.path.join(self._get_task_file_dir(task), "steering_file_wrapper.py")
        gb2_mock_process.pickle_file_path = self.test_dir
        gb2_mock_process.wrapper_file_path = self.test_dir
        gb2_mock_process.gbasf2_project_name = get_unique_project_name(task)
        gb2_mock_process.task = task
        return Gbasf2Process._build_gbasf2_submit_command(gb2_mock_process)

    def test_input_dataset_in_command(self):
        "Test that if the input dataset is set, it occurs once in the command string"

        class _Task(MyGbasf2Task):
            gbasf2_input_dataset = b2luigi.Parameter(default=self.dummy_lfn)

        gb2_cmd = self._build_gbasf2_submit_command(_Task("some_parameter"))
        gb2_cmd_str = " ".join(gb2_cmd)
        self.assertEqual(gb2_cmd.count(self.dummy_lfn), 1)
        self.assertEqual(gb2_cmd.count("-i"), 1)
        self.assertEqual(gb2_cmd_str.count(f" -i {self.dummy_lfn}"), 1)
        self.assertNotIn("--input_dslist", gb2_cmd_str)

    def test_input_dslist_in_command(self):
        "Test that if the input dataset list is set, it occurs once in the command string"

        class _Task(MyGbasf2Task):
            gbasf2_input_dslist = self.dummy_lfn_file.name

        gb2_cmd = self._build_gbasf2_submit_command(_Task("some_parameter"))
        gb2_cmd_str = " ".join(gb2_cmd)
        self.assertEqual(gb2_cmd.count(self.dummy_lfn_file.name), 1)
        self.assertEqual(gb2_cmd.count("--input_dslist"), 1)
        self.assertEqual(gb2_cmd_str.count(f" --input_dslist {self.dummy_lfn_file.name}"), 1)
        self.assertNotIn("-i", gb2_cmd)

    def test_input_dslist_set_but_does_not_exist(self):
        """Test that if the input dataset list is set but the file does not exist, an ``FileNotFoundError`` is raised"""

        class _Task(MyGbasf2Task):
            gbasf2_input_dslist = self.dummy_lfn_file.name

        self.dummy_lfn_file.close()
        self.assertRaises(FileNotFoundError, lambda: self._build_gbasf2_submit_command(_Task("some_parameter")))

    def test_not_dataset_raises_error(self):
        """Test that we get an error when no gbasf2 input dataset is provided"""
        self.assertRaises(RuntimeError, lambda: self._build_gbasf2_submit_command(MyGbasf2Task("some_parameter")))

    def test_both_input_dataset_and_dslist_raises_error(self):
        """Test that en error is raise when both gbasf2_input_dataset and gbasf2_inputdslist settings are given"""

        class _Task(MyGbasf2Task):
            gbasf2_input_dataset = b2luigi.Parameter(default=self.dummy_lfn)
            gbasf2_input_dslist = self.dummy_lfn_file.name

        self.assertRaises(RuntimeError, lambda: self._build_gbasf2_submit_command(MyGbasf2Task("some_parameter")))


class TestGetGbasf2DatasetQuery(unittest.TestCase):
    "Run tests for ``Gbasf2Process._get_gbasf2_dataset_query``."

    def setUp(self):
        super().setUp()
        self.gb2_mock_process = Mock()
        self.gb2_mock_process.gbasf2_project_name = "projectname"
        self.gb2_mock_process.dirac_user = "username"
        self.gb2_mock_process.task = MyTask("some_parameter")

    def test_filename_single_root_extension(self):
        dataset_query = Gbasf2Process._get_gbasf2_dataset_query(self.gb2_mock_process, output_file_name="output.root")
        self.assertEqual(dataset_query, "/belle/user/username/projectname/sub*/output_*.root")

    def test_filename_multi_extension(self):
        dataset_query = Gbasf2Process._get_gbasf2_dataset_query(
            self.gb2_mock_process, output_file_name="output.mdst.root"
        )
        self.assertEqual(dataset_query, "/belle/user/username/projectname/sub*/output_*.mdst.root")

    def test_non_root_extension_raises_err(self):
        "Raise an error if gbasf2 output does not end with .root"
        with self.assertRaises(ValueError):
            Gbasf2Process._get_gbasf2_dataset_query(self.gb2_mock_process, output_file_name="output.mdst")

    def test_non_basename_raises_err(self):
        """
        Currently we only support simple basenames as gbasf2 job outputs.
        LPN is guessed via dirac username for user projects.
        """
        with self.assertRaises(ValueError):
            Gbasf2Process._get_gbasf2_dataset_query(self.gb2_mock_process, output_file_name="/path/to/output.root")

    def test_filename_single_extension(self):
        gb2_mock_process = Mock()
        gb2_mock_process.gbasf2_project_name = "projectname"
        gb2_mock_process.dirac_user = "username"
        gb2_mock_process.task = MyTask("some_parameter")

        dataset_query = Gbasf2Process._get_gbasf2_dataset_query(self=gb2_mock_process, output_file_name="output.root")
        self.assertEqual(dataset_query, "/belle/user/username/projectname/sub*/output_*.root")
