#!/usr/bin/env python
"""
Copy EXR file from src to destination and set additional
render related attributes.

This script is based on Sony's original version.

Since PRMan 17 supports arbitrary meta data, we can no longer have a fixed list
of attributes we want to copy. We exclude certain attributes that previously 
were not copied. Furthermore, only some types are supported by OpenEXR and
PyOpenEXR. For those that are not, the script will copy the attributes as
strings, e.g. float[4] which is box2f. 'box2f' is not supported by the 
PyOpenEXR.WriteExrFile function. This will hopefully be fixed in a future 
version.
"""
import os
import tempfile
import shutil
import sys
from optparse import OptionParser
import PyOpenEXR as OpenEXR 
import PyImath as Imath

# These attributes do not need to be copied because they already exist in the
# destination file.
EXCLUDE_ATTRIBUTES = ('dataWindow', 'displayWindow','channels','compression', 
                      'lineOrder', 'tiles')

# See the bottom of 'openexr-1.6.1/IlmImf/ImfHeader.cpp' for all types. Types  
# like box2f and box2i are mentioned there, but not supported by 
# PyOpenEXR.WriteExrFile.
SUPPORTED_ATTRIBUTE_TYPES = ["int", "float", "string", "double", 
                             "m33f", "m44f", "v2f", "v2i", "v3f", "v3i"]

def main():
    """
    Parse command line options.
    """
    usage = "Usage: %prog [options] SRC_EXR DEST_EXR"
    parser = OptionParser(usage=usage)
    parser.add_option('--package', dest='package',
                      help='String to add to header as "package" attr.')
    parser.add_option('--renderer', dest='renderer',
                      help='String to add to header as "renderer" attr.')
    parser.add_option('--attr', dest='attrs', action='append',
                      help='Arbitrary attribute name and value.')
    parser.add_option('--verbose', dest='verbose',
                      help='Outputs debug information about which '
                      'attributes have been copied or skipped.')

    (options, args) = parser.parse_args()

    if len(args) != 2:
        parser.error('incorrect number of EXR file inputs, %s' % usage)

    sourceAttrs = OpenEXR.GetExrAttributes(args[0])
    exrAttrs = { }

    srcFile = args[0]
    dstFile = args[1]

    # Copy attrs from source EXR that are supported by PyOpenEXR
    for attrName in sourceAttrs:
        attrType = sourceAttrs[attrName][0]
        if attrType in SUPPORTED_ATTRIBUTE_TYPES:
            value = sourceAttrs[attrName][1]
            # Oddly enough, PyOpenEXR needs as to pass a float for int attrs
            if attrType == 'int':
                value = float(value)
            exrAttrs[attrName] = (sourceAttrs[attrName][0], value)
            
            if options.verbose:
                print("Copied '%s' of type '%s'."
                      % (attrName, sourceAttrs[attrName][0]))
        elif attrName not in EXCLUDE_ATTRIBUTES:
            # Copy data as string for unsupported types.
            exrAttrs[attrName] = ('string', str(sourceAttrs[attrName][1]))
            
            if options.verbose:
                print("Type '%s' is unsupported by exrcopyrenderattrs script. "
                      "Added '%s' as string."
                      % (sourceAttrs[attrName][0], attrName))
        elif options.verbose:            
            print("Skipped copying of '%s' attribute." % (attrName))

    if options.package:
        exrAttrs['package'] = ('string', options.package)

    if options.renderer:
        exrAttrs['renderer'] = ('string', options.renderer)

    if options.attrs:
        try:
            handleAttrFlag(options.attrs, exrAttrs)
        except Exception, e:
            sys.stderr.write("Failed to handle the --attr flag, %s\n" % e)

    exrCopy(srcFile, dstFile, exrAttrs)

def handleAttrFlag(attrs, exrAttrs):
    """
    Handle --attr flags.
    """
    for attr in attrs:
        key, val = attr.split("=", 1)
        try:
            value = float(val)
            exrAttrs[key] = ('float', value)
        except ValueError:
            exrAttrs[key] = ('string', val)

def exrCopy(src, dst, attrs):
    """
    Copy EXR from source to destination.
    """
    if not OpenEXR.IsExrFile(src):
        sys.exit("Source file '%s' is not an exr." % src)

    if not OpenEXR.IsExrFile(dst):
        sys.exit("Destination file '%s' is not an exr." % dst)

    # Create a safe tempfile
    (tmpfd, tmpfilename) = tempfile.mkstemp('.exr', '', 
                                            os.environ.get('TMPDIR',
                                                tempfile.gettempdir()),
                                            True)

    try:
        OpenEXR.WriteExrFile(tmpfilename, dst, attrDict=attrs)
        shutil.copyfile(tmpfilename, dst)
    except Exception, e:
        msg = "Error writing EXR %s, %s.\n" % (dst, e)
        sys.stderr.write(msg)
        sys.exit(1)
    finally:
        try:
            os.close(tmpfd)
            os.remove(tmpfilename)
        except:
            pass

if __name__ == '__main__':
    main()
