# =============================================================================
#
#  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.model
import smtk.operation
import smtk.project
import smtk.resource
import smtk.view

import smtk.testing

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

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 = '52218014'
    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)
        self.proj_manager.registerProject('basic', set(), set(), '0.0')

        # 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)."""
        project = self.proj_manager.createProject('smtk::simulation::ace3p::Project')
        self.assertIsNotNone(project, 'failed to create project')

        # Set the project name and location
        project_name = 'test'
        project.setName(project_name);
        smtk_path =  '{}/{}.project.smtk'.format(self.project_folder, project_name)
        project.setLocation(smtk_path);
        print('USING SMTK_PATH', smtk_path)

        # Create attribute resource
        att_resource = self.res_manager.createResource('smtk::attribute::Resource')
        self.res_manager.add(att_resource)
        self.assertIsNotNone(att_resource)
        sbt_path = os.path.join(smtk.testing.SOURCE_DIR, 'simulation-workflows/ACE3P.sbt')
        import_op = self.op_manager.createOperation('smtk::attribute::Import')
        import_op.parameters().associate(att_resource)
        import_op.parameters().findFile('filename').setValue(sbt_path)
        import_result = import_op.operate()
        outcome = import_result.findInt('outcome').value()
        self.assertEqual(outcome, OP_SUCCEEDED, 'Import sbt failed, outcome {}'.format(outcome))

        # Manually create analysis attribute
        view = att_resource.findViewByType('Analysis')
        self.assertIsNotNone(view, 'attribute template missing Analysis view')
        view_atts = view.details().attributes()
        att_name = view_atts.get('AnalysisAttributeName')
        defn_type = view_atts.get('AnalysisAttributeType')
        defn = att_resource.analyses().buildAnalysesDefinition(att_resource, defn_type)
        self.assertIsNotNone(defn, 'missing analysis definition, type {}'.format(defn_type))
        att = att_resource.createAttribute(att_name, defn)
        self.assertIsNotNone(att, 'failed to create analysis attribute')

        # Add attribute resource to the project
        added = project.resources().add(att_resource, 'attributes')
        self.assertTrue(added, 'failed to add attribute resource to project')
        self.res_manager.add(att_resource)
        self.assertEqual(project.resources().size(), 1, \
            'Wrong number of project resources; should be 1 not {}'.format(project.resources().size()))

        # Add stage to project
        model_resource = smtk.model.Resource.create() # not used
        stage = smtk.simulation.ace3p.Stage(att_resource, model_resource, project.jobsManifest(), 'Stage1')
        project.addStage(stage)

        # Write project
        write_op = self.op_manager.createOperation('smtk::simulation::ace3p::Write')
        self.assertIsNotNone(write_op)
        write_op.parameters().associate(project)
        write_result = write_op.operate()
        outcome = write_result.findInt('outcome').value()
        self.assertEqual(outcome, OP_SUCCEEDED, 'Write project failed, outcome {}'.format(outcome))

        return project

    def init_export_params(self, params, project):
        """Initialize baseline export parameters"""
        params.associate(project)
        params.findInt('stage-index').setValue(0)

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

        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')

    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)

        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)

        """Skip following checks because logic to construct job record has moved from
           Project.cxx to qtNewtJobSubmitter, which is not accessible from python.
        """

        # 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()
