# =============================================================================
#
#  Copyright (c) Kitware, Inc.
#  All rights reserved.
#  See LICENSE.txt for details.
#
#  This software is distributed WITHOUT ANY WARRANTY; without even
#  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
#  PURPOSE.  See the above copyright notice for more information.
#
# =============================================================================
"""Test Project::toJobRecord()."""

import os
import shutil
import sys

import smtk
import smtk.attribute
import smtk.io
import smtk.operation
import smtk.project
import smtk.resource

try:
    import smtksimulationace3p # configures smtk.simulation.ace3p
except ImportError:
    print('Failed to import smtksimulationace3p. sys.path:')
    for p in sys.path:
        print('  ', p)
        raise

import smtk.testing

OP_SUCCEEDED = int(smtk.operation.Operation.Outcome.SUCCEEDED)

class TestData:
    """Data to put into export objects and end up in job record.

    Some of the values are initialized in the test code
    """
    cumulus_id = '60c3a9e8892404f436fc5754'
    job_name = 'example1-acdtool'
    machine = 'cori'
    analysis = 'ACDTool'
    analysis_id = None  # initialized in init_export_params()
    nodes = 16
    processes = 128
    runtime = None  # not tested
    submission_time = None  # not tested
    notes = 'the quick brown fox jumps over the lazy dog'
    runtime_job_folder = '/global/cscratch1/sd/johnt/210611/example1-acdtool/60c3a9e8892404f436fc5754'
    local_job_folder = None  # initialized in setUp()
    runtime_input_folder = '/global/cscratch1/sd/johnt/test/ACE3P/60b10729892404f436fc5480'


class TestToJobRecord(smtk.testing.TestCase):
    def setUp(self):
        # Set project folder
        self.project_folder = os.path.join(smtk.testing.TEMP_DIR, 'python/test_to_jobrecord')
        if os.path.exists(self.project_folder):
            shutil.rmtree(self.project_folder)
        TestData.local_job_folder = os.path.join(self.project_folder, 'jobs', TestData.cumulus_id)

        # Initialize smtk managers
        self.res_manager = smtk.resource.Manager.create()
        self.op_manager = smtk.operation.Manager.create()

        smtk.attribute.Registrar.registerTo(self.res_manager)
        smtk.attribute.Registrar.registerTo(self.op_manager)

        smtk.operation.Registrar.registerTo(self.op_manager)
        self.op_manager.registerResourceManager(self.res_manager)

        self.proj_manager = smtk.project.Manager.create(self.res_manager, self.op_manager)
        smtk.simulation.ace3p.Registrar.registerTo(self.proj_manager)

        # Set workflows directory
        workflows_dir = os.path.join(smtk.testing.SOURCE_DIR, 'simulation-workflows')
        smtk.simulation.ace3p.Metadata.WORKFLOWS_DIRECTORY = workflows_dir

    def tearDown(self):
        self.res_manager = None
        self.op_manager = None
        self.proj_manager = None

    def create_project(self):
        """Create ACE3P project with attribute resource only (no model)."""
        create_op = self.op_manager.createOperation('smtk::simulation::ace3p::Create')
        self.assertIsNotNone(create_op)
        create_op.parameters().findDirectory("location").setValue(self.project_folder)
        create_op.parameters().find('overwrite').setIsEnabled(True)
        create_op.parameters().find('analysis-mesh').setIsEnabled(False)
        self.assertTrue(create_op.ableToOperate())

        result = create_op.operate()
        outcome = result.findInt('outcome').value()
        if outcome != OP_SUCCEEDED:
            print(create_op.log().convertToString())
        self.assertEqual(outcome, OP_SUCCEEDED)
        project = result.findResource('resource').value()
        self.assertIsNotNone(project)
        return project

    def init_export_params(self, params, project):
        """Initialize baseline export parameters."""
        # Get the attribute resource and assign the "analysis" item
        resource_set = project.resources().find('smtk::attribute::Resource')
        att_resource = resource_set.pop()

        # Set analysis type
        analysis_att = att_resource.findAttribute('analysis')
        analysis_item = analysis_att.findString('Analysis')
        isset = analysis_item.setValue(TestData.analysis)
        self.assertTrue(isset)

        params.findResource('analysis').setValue(att_resource)
        TestData.analysis_id = str(att_resource.id())

        output_folder = os.path.join(self.project_folder, 'export')
        params.findDirectory('OutputFolder').setValue(output_folder)

        params.findString('OutputFilePrefix').setValue('jobrecordtest')

        # Hack in mesh file (without loading)
        meshfile_path = os.path.join(smtk.testing.DATA_DIR, 'models/3d/genesis/pillbox4.gen')
        params.findFile('MeshFile').setValue(meshfile_path)

    def init_submit_params(self, params):
        """"""
        sim_item = params.findGroup('NERSCSimulation')
        sim_item.setIsEnabled(True)
        sim_item.find('JobName').setValue(TestData.job_name)
        sim_item.find('JobNotes').setValue(TestData.notes)
        sim_item.find('Machine').setValue(TestData.machine)
        sim_item.find('NumberOfNodes').setValue(TestData.nodes)
        sim_item.find('NumberOfTasks').setValue(TestData.processes)

    def init_result(self, result_att):
        """"""
        # Create fake export file
        export_folder = os.path.join(self.project_folder, 'export')
        os.makedirs(export_folder)
        output_file = os.path.join(export_folder, 'ace3p.rfpost')
        with open(output_file, 'w') as out:
            out.write('xyzzy\n')
            print('Wrote file', output_file)

        result_att.findFile('OutputFile').setIsEnabled(True)
        result_att.findFile('OutputFile').setValue(output_file)
        result_att.findString('CumulusJobId').setIsEnabled(True)
        result_att.findString('CumulusJobId').setValue(TestData.cumulus_id)
        result_att.findString('NerscJobFolder').setIsEnabled(True)
        result_att.findString('NerscJobFolder').setValue(TestData.runtime_job_folder)
        result_att.findString('NerscInputFolder').setIsEnabled(True)
        result_att.findString('NerscInputFolder').setValue(TestData.runtime_input_folder)

    def check_jobrecord(self, project):
        """"""
        self.assertEqual(project.jobData(0, 'cumulus_id'), TestData.cumulus_id)
        self.assertEqual(project.jobData(0, 'job_name'), TestData.job_name)
        self.assertEqual(project.jobData(0, 'machine'), TestData.machine)
        self.assertEqual(project.jobData(0, 'analysis'), TestData.analysis)
        self.assertEqual(project.jobData(0, 'analysis_id'), TestData.analysis_id)
        self.assertEqual(int(project.jobData(0, 'nodes')), TestData.nodes)
        self.assertEqual(int(project.jobData(0, 'processes')), TestData.processes)
        self.assertEqual(project.jobData(0, 'notes'), TestData.notes)
        self.assertEqual(project.jobData(0, 'runtime_job_folder'), TestData.runtime_job_folder)
        self.assertEqual(project.jobData(0, 'local_job_folder'), TestData.local_job_folder)
        self.assertEqual(project.jobData(0, 'runtime_input_folder'), TestData.runtime_input_folder)

    def test_to_jobrecord(self):
        """"""
        project = self.create_project()

        # Construct export spec including result attribute
        export_op = self.op_manager.createOperation('smtk::simulation::ace3p::Export')
        self.assertIsNotNone(export_op)
        export_op.parameters().associate(project)

        params = export_op.parameters()
        params.find('test-mode').setIsEnabled(True)
        self.init_export_params(params, project)
        self.init_submit_params(params)
        result = export_op.createResult(smtk.operation.Operation.Outcome.SUCCEEDED)
        self.init_result(result)

        # Add job record to project
        logger = smtk.io.Logger.instance()
        logger.reset()
        added = project.onJobSubmit(params, result, logger)
        if not added:
            print(logger.convertToString())
        if logger.numberOfRecords():
            print('LOGGER: ', logger.convertToString())
        self.assertTrue(added)

        # Finally we can check the job record
        self.check_jobrecord(project)

        # Finis
        shutil.rmtree(self.project_folder)


if __name__ == '__main__':
    smtk.testing.process_arguments()
    smtk.testing.main()
