# =============================================================================
#
#  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.
#
# =============================================================================
"""ImportModel operation"""

import os
import shutil
import sys

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

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

# Make sure __file__ is set when using modelbuilder
import inspect
source_file = os.path.abspath(inspect.getfile(inspect.currentframe()))
__file__ = source_file

# Make sure this folder is in sys.path
source_dir = os.path.abspath(os.path.dirname(__file__))
if source_dir not in sys.path:
    sys.path.insert(0, source_dir)

from foam_mixin import FoamMixin


class ImportModel(smtk.operation.Operation, FoamMixin):
    """Operation to import geometry file into wind tunnel project"""

    def __init__(self):
        smtk.operation.Operation.__init__(self)
        FoamMixin.__init__(self)
        # Do NOT store any smtk resources as member data (causes memory leak)

    def name(self):
        return "Import geometry file into wind tunnel project"

    def createSpecification(self):
        spec = self._create_specification(sbt_file='import_model.sbt')
        return spec

    def operateInternal(self):
        """"""
        project = self._get_project()
        if project is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        att_resource = self._get_attribute_resource(project)
        if att_resource is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        proj_manager = project.manager()
        res_manager = proj_manager.resourceManager()
        op_manager = proj_manager.operationManager()

        project_dir = os.path.abspath(os.path.dirname(project.location()))

        # Make sure assets directory exists
        assets_dir = os.path.join(project_dir, 'assets')
        if not os.path.exists(assets_dir):
            os.makedirs(assets_dir)

        # Get attribute resource and 'Target' attribute
        target_att = att_resource.findAttribute('Target')
        target_filename_item = target_att.findString('Filename')

        # Check for existing model
        model_resource_set = project.resources().findByRole('model')
        if model_resource_set:
            model_resource = model_resource_set.pop()
            overwrite_item = self.parameters().findVoid('overwrite')
            if overwrite_item.isEnabled():
                att_resource.disassociate(model_resource)

                project.resources().remove(model_resource)
                res_manager.remove(model_resource)

                # Delete the geometry file from the assets folder
                target_filename = target_filename_item.value()
                target_path = os.path.join(assets_dir, target_filename)
                if os.path.exists(target_path):
                    os.remove(target_path)

                # If there is .vtp file, delete that too
                basename, _ = os.path.splitext(target_filename)
                vtp_filename = '{}.vtp'.format(basename)
                vtp_path = os.path.join(assets_dir, vtp_filename)
                if os.path.exists(vtp_path):
                    os.remove(vtp_path)

                # Clear Filename item
                target_filename_item.setValue('')
                target_filename_item.setIsEnabled(False)
            else:
                self.log().addError('model already exists and overwrite flag is false')
                return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Copy geometry file to project assets directory
        input_path = self.parameters().findFile('filename').value()
        input_filename = os.path.basename(input_path)
        asset_path = os.path.join(assets_dir, input_filename)
        shutil.copyfile(input_path, asset_path)

        # Check option to import model as vtp file.
        # Experience is that the smtk mesh resource does not reliably import
        # some .stl or .obj files, but *can* import the same geometry when
        # converted to .vtp (vtkPolyData) format. When the 'import-vtp' option
        # is enabled, the input file will be (i) copied to the project assets
        # directory and (ii) converted to the equivalent .vtp file and that
        # is also stored in the project assets directory.
        # The .vtp file will then be imported as the model resource, and the
        # original/native file will be used for meshing.
        # This option is enabled by default.
        import_path = asset_path
        vtp_item = self.parameters().findVoid('import-vtp')
        if vtp_item is not None and vtp_item.isEnabled():
            vtp_path = self._convert_to_vtp(input_path, assets_dir)
            if vtp_path is not None:
                import_path = vtp_path

        # Import geometry file
        # For expediency, load as mesh resource (smtk::mesh::Resource)
        # instead of model resource, because default modelbuilder does not
        # package mesh *session* (which can read/write model resource represented
        # by smtk mesh.
        import_op = op_manager.createOperation('smtk::operation::ImportResource')
        import_op.parameters().findFile('filename').setValue(import_path)
        import_result = import_op.operate()
        import_outcome = import_result.findInt('outcome').value()
        if import_outcome != OP_SUCCEEDED:
            self.log().addError('Import file operation returned outcome {}'.format(import_outcome))
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        model_resource = import_result.findResource('resourcesCreated').value()
        if model_resource is None:
            self.log().addError('Model resource not returned by smtk::operation::ImportResource')
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        project.resources().add(model_resource, 'model')

        # Update attribute resource
        att_resource.associate(model_resource)
        target_filename_item.setIsEnabled(True)
        target_filename_item.setValue(input_filename)

        # Update analysis items to set the "target" category
        # This logic is not typical for applications, but enables us to use categories to
        # both: (i) set the solver options, and (ii) enable views that only apply to the
        # imported model. For reference, see the Analysis definition in WindTunnel.sbt
        # Future: This logic should probably be migrated to some common/shared source.
        analysis_att = att_resource.findAttribute('Analysis')
        analysis_item = analysis_att.findString('Analysis')
        item_dict = analysis_item.childrenItems()
        for name, group_item in item_dict.items():
            target_item = group_item.find('target', smtk.attribute.SearchStyle.IMMEDIATE)
            target_item.setIsEnabled(True)
        item_dict = None
        analysis_item = None
        analysis_att = None

        # Add 'target' to active categories
        cats = att_resource.activeCategories()
        cats.add('target')
        att_resource.setActiveCategories(cats)

        # Write project to disk
        write_op = self.manager().createOperation('smtk::project::Write')
        write_op.parameters().associate(project)
        write_result = write_op.operate()
        write_outcome = write_result.findInt('outcome').value()
        if (write_outcome != OP_SUCCEEDED):
            self.log().addError('Write project failed, outcome {}'.format(write_outcome))
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        result = self.createResult(smtk.operation.Operation.Outcome.SUCCEEDED)
        # Add model resource to results (attribute)
        resource_item = result.findResource('resourcesCreated')
        ok = resource_item.setNumberOfValues(1)
        ok &= resource_item.setValue(model_resource)
        if not ok:
            self.log().addError('Failed to assign new resource to results item')
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Add analysis attribute to results 'modified' item
        result.findComponent('modified').appendValue(analysis_att)
        # Add target attribute to results 'modified' item
        result.findComponent('modified').appendValue(target_att)
        return result

    def _convert_to_vtp(self, input_path: str, assets_dir: str) -> str:
        """Converts input file to vtkPolyData and stores vtp file.

        param input_path: path to input geometry file
        param assets_dir: path to target directory for .vtp file
        :return: path to .vtp file in the assets directory or None if not converted.
        """
        input_filename = os.path.basename(input_path)
        basename, ext = os.path.splitext(input_filename)

        # Import vtk modules and create reader
        from vtkmodules.vtkCommonDataModel import vtkPolyData
        from vtkmodules.vtkIOXML import vtkXMLPolyDataWriter
        if ext == '.obj':
            from vtkmodules.vtkIOGeometry import vtkOBJReader
            reader = vtkOBJReader()
        elif ext in ['.stl', '.stlb']:
            from vtkmodules.vtkIOGeometry import vtkSTLReader
            reader = vtkSTLReader()
        else:
            self.log().addError('Unsupported file extension {}'.format(ext))
            return None

        reader.SetFileName(input_path)
        reader.Update()
        polydata = reader.GetOutput()

        output_filename = '{}.vtp'.format(basename)
        output_path = os.path.join(assets_dir, output_filename)

        writer = vtkXMLPolyDataWriter()
        writer.SetFileName(output_path)
        writer.SetDataModeToAscii()
        # writer.SetInputConnection(reader.GetOutputPort())
        writer.SetInputData(polydata)
        writer.Write()

        return output_path
