#=============================================================================
#
#  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.
#
#=============================================================================
"""AttributeBuilder script that uses smtk-tools submodule."""

import argparse
import os
import sys
sys.dont_write_bytecode = True

# Make sure we can import all smtk modules
try:
    import smtk
    import smtk.attribute
    import smtk.io
    import smtk.model
    import smtk.operation
    import smtk.session.vtk

except ImportError:
    print('sys.path:')
    for p in sys.path:
        print('  ', p)
    print()
    print('ERROR: FAILED TO IMPORT SMTK MODULES')
    print('You probably need to set PYTHONPATH to <smtk-build-directory>/lib/python3.9/site-packages')
    sys.exit(-1)


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


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Populate project attribute resource (one stage) from yml description')
    # Required arguments
    parser.add_argument('project_directory', help='Project directory')
    parser.add_argument('yml_file', help='YAML file specifying the attributes to generate')
    if sys.platform == 'win32':
        parser.add_argument('ace3p_dll_directory', help='Path to ace3p-extensions/bin directory')

    # Optional arguments
    parser.add_argument('-s', '--stage_index', type=int, default=0, help='stage number')
    parser.add_argument('-t', '--smtktools_path', help='path to smtk_tools module')
    # parser.add_argument('-d', '--dry-run', action='store_true', help='dont write result to file system')

    args = parser.parse_args()
    # print(args)

    # Need to set dll directory *before* import smtksimulationace3p
    if 'ace3p_dll_directory' in args:
        if not os.add_dll_directory(args.ace3p_dll_directory):
            print('ERROR: failed to add ace3p_dll_directory {}'.format(args.ace3p_dll_directory))
            sys.exit(-1)

    try:
        import smtksimulationace3p
    except ImportError:
        print('ERROR: FAILED TO IMPORT ACE3P MODULE (smtksimulationace3p)')
        if sys.platform == 'win32':
            print('Check that your ace3p_dll_directory argument is correct')
        sys.exit(-1)

    # Check filesystem paths
    assert os.path.exists(args.project_directory), \
        'project_directory not found: {}'.format(args.project_directory)
    project_filename = '{}.project.smtk'.format(os.path.basename(args.project_directory))
    project_filepath = os.path.join(args.project_directory, project_filename)
    assert os.path.exists(project_filepath), \
        'project resource file not found: {}'.format(project_filepath)
    assert os.path.exists(args.yml_file), \
        'yml spec not found: {}'.format(args.yml_file)
    if args.smtktools_path is not None and not os.path.exists(args.smtktools_path):
        raise RuntimeError('smtk_tools directory not found: {}'.format(args.smtktools_path))

    # Add path to smtk_tools
    tools_path = args.smtktools_path
    if tools_path is None:
        # Default path from the repository
        source_path = os.path.abspath(os.path.dirname(__file__))
        pardirs = [os.pardir]*5
        path = os.path.join(source_path, *pardirs, 'thirdparty', 'smtk-tools')
        tools_path = os.path.normpath(path)

    assert os.path.exists(tools_path), 'smtk-tools path not found at {}'.format(tools_path)
    sys.path.insert(0, tools_path)

    import smtk_tools
    from smtk_tools.resource_io import ResourceIO
    from smtk_tools.attribute_builder import AttributeBuilder

    # Load yml file as the attribute specification
    import yaml
    spec = None
    print('Loading yaml file:', args.yml_file)
    with open(args.yml_file) as fp:
        content = fp.read()
        spec = yaml.safe_load(content)
    assert spec is not None, 'failed to load/parse yml file'

    # Load project
    print('Loading project resource:', project_filepath)
    res_io = ResourceIO()
    # Hack loading project pending update to smtk-tools
    proj_manager = smtk.project.Manager.create(res_io.res_manager, res_io.op_manager)
    smtk.simulation.ace3p.Registrar.registerTo(proj_manager)

    # project = smtk.simulation.ace3p.Project.create()
    read_op = res_io.op_manager.createOperation('smtk::simulation::ace3p::Read')
    assert read_op is not None
    read_op.parameters().findFile('filename').setValue(project_filepath)
    read_result = read_op.operate()
    read_outcome = read_result.findInt('outcome').value()
    assert read_outcome == OP_SUCCEEDED
    project = read_result.findResource('resource').value()
    # print('project:', project)

    # Get the attribute resource
    print('Number of stages:', project.numberOfStages())
    print('Using project stage {}'.format(args.stage_index))
    stage = project.stage(args.stage_index)
    assert stage is not None
    att_resource = stage.attributeResource()
    model_resource = stage.modelResource()

    # Apply the spec
    print('Building attributes')
    builder = AttributeBuilder()
    builder.build_attributes(att_resource, spec, model_resource=model_resource, create_instances=False)
    att_resource.setClean(False)

    # Write the project, again hacking ResourceIO
    print('Writing project')
    write_op = res_io.op_manager.createOperation('smtk::simulation::ace3p::Write')
    assert write_op is not None
    write_op.parameters().associate(project)
    write_result = write_op.operate()
    write_outcome = write_result.findInt('outcome').value()
    assert write_outcome == OP_SUCCEEDED, 'failed to write project resource'

    print('finis')
