#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-
#*****************************************************************************/
# Copyright (c) 2006 Autodesk, Inc.
# All rights reserved.
#
# These coded instructions, statements, and computer programs contain
# unpublished proprietary information written by Autodesk, Inc., and are
# protected by Federal copyright law. They may not be disclosed to third
# parties or copied or duplicated in any form, in whole or in part, without
# the prior written consent of Autodesk, Inc.
#*****************************************************************************/

"""
Import 1rst party render output in Toxik

Usage:

  toxik-imsq-import.py <imsq-filename>

This script creates a new composition (or reuses an existing one) in
the Toxik database to correspond to Maya or 3dsmax files to be
rendered later.  It reads the list of files and the other parameters
from the IMSQ file.

Note: this script is meant to be invoked from Maya or Max that is
triggered before rendering the frames out.  """

#==============================================================================
# EXTERNAL DECLARATIONS
#==============================================================================

import sys, os, os.path, re, fnmatch, posixpath, datetime, math
from os.path import *
from xml.dom.minidom import parse as domparse

linux = (os.name == 'posix')

blendMap = {
  'Normal' : 0,
  'Average' : 1,
  'Add' : 2,
  'Substract' : 3,
  'Darken' : 4,
  'Multiply' : 5,
  'Color Burn' : 6,
  'Linear Burn' : 7,
  'Lighten' : 8,
  'Screen' : 9,
  'Color Dodge' : 10,
  'Linear Dodge' : 11,
  'Spotlight' : 12,
  'Spotlight Blend' : 13,
  'Overlay' : 14,
  'Soft Light' : 15,
  'Hard Light' : 16,
  'Pin Light' : 17,
  'Hard Mix' : 18,
  'Difference' : 19,
  'Exclusion' : 20,
  'Hue' : 21,
  'Saturation' : 22,
  'Color' : 23,
  'Value' : 24
}

#==============================================================================
# LOCAL DECLARATIONS
#==============================================================================

searchPattern = re.compile('\.#*\.')

def toxik_environment_fixup(opts):
    """
    This script is meant to be run directly from Max or Maya, for Toxik.
    This more or less fixes up the environment to be able to load Toxik modules.
    """

    # These need to be changed for your needs
    if opts.toxikPath:
        install_dir = opts.toxikPath
    else:
        if linux:
            install_dir = '/opt/Autodesk/Autodesk_Toxik-2007'
        else:
            install_dir = 'C:\\Program Files\\Autodesk\\Autodesk Toxik 2007'

    os.environ['DL_INSTALL_DIR'] = install_dir
    os.environ['DL_PLUGINS_DIR'] = os.path.join(install_dir, 'plugins')
    if opts.verbose:
        print >> sys.stderr, "DL_INSTALL_DIR =", os.environ['DL_INSTALL_DIR']
        print >> sys.stderr, "DL_PLUGINS_DIR =", os.environ['DL_PLUGINS_DIR']

    # Add DLL path so that we can find the program DLLs when the toxik modules
    # and/or plugins are loaded in memory.
    if not linux:
        os.environ['PATH'] = '%s;%s' % (os.path.join(install_dir, 'program'),
                                        os.environ['PATH'])
        if opts.verbose:
            print >> sys.stderr, 'PATH =', os.environ['PATH']

    # Set PYTHONPATH so that we can find the Toxik modules.
    sys.path.append(os.path.join(install_dir, 'python', 'lib', 'python2.4'))
    if opts.verbose:
        print >> sys.stderr, 'PYTHONPATH =', sys.path

    # Set the media cache to a new place due to a concurrency problem with Toxik
    # and its media cache. Note: the directory must exist.
    if opts.tempFolder:
        media_cache_root = opts.tempFolder
        if not os.path.isdir(media_cache_root):
            raise RuntimeError("Media cache folder specified by --tempFolder "
                               "argument is not a directory: %s" % media_cache_root)
    else:
        import tempfile        
        media_cache_root = tempfile.gettempdir()

    for x in ('store', 'index'):
        mcdir = os.path.join(media_cache_root, x)
        os.environ['DL_MEDIA_CACHE_%s' % x.upper()] = mcdir
        if not os.path.exists(mcdir):
            os.mkdir(mcdir)

#------------------------------------------------------------------------------
#
def genFileList(file_pattern, frame_start, frame_end):
    """
    Expands the pattern and frame numbers to a list of filenames.
    The end frame index is exclusive at this point.
    The pattern is expected to contain a marker like [#...] to be replaced, if
    there are more than one frame to generate.

    An example:

      >>> genFileList('bl[##]', 2, 6)
      ['bl02', 'bl03', 'bl04', 'bl05']

      >>> genFileList('bl[##]', 0, 1)
      ['bl00']

      >>> genFileList('bl', 0, 1)
      ['bl']
    """
    assert frame_end >= frame_start
    nframes = frame_end - frame_start
    if nframes == 0:
        return []

    mo = searchPattern.search(file_pattern)
    if mo:
        assert nframes >= 1
        nchars = len(mo.group()) - 2
        rep_pattern = searchPattern.sub('.%%0%dd.' % nchars, file_pattern)
        return map(lambda x: rep_pattern % x, xrange(frame_start, frame_end))
    else:
        assert nframes == 1
        return [file_pattern]

#------------------------------------------------------------------------------
#
def createComps(dbname, project, username, footageRoot, compName, rate, 
    baseFolder, sequences, opts):
    """
    Given a database name and project name, a file pattern (string) and a start
    end end frame indexes (ints), create or get a composition and add a new
    rendered result to hold the frames specified in the pattern and indices.

    We expect the 'compname' to be a database path relative to its project
    folder and not to the root of the database.
    """

    # fixup the environment before importing the Toxik modules
    toxik_environment_fixup(opts)
    
    if not opts.test:
        from autodesk_toxik import database, pimgr, admin, graph, media
        if opts.verbose:
            print >> sys.stderr, 'Toxik modules imported correctly'
    else:
        graph = toxik_Stub()
        database = toxik_Stub()
        pimgr = toxik_Stub()
        admin = toxik_Stub()
        media = toxik_Stub()
        if opts.verbose:
            print >> sys.stderr, 'Using Toxik stubs'
        
    # Note: we need to get the following render parameters from the file:
    # - height,
    # - width,
    # - rate,
    # - components

    # connect to the database
    if opts.verbose:
        print >> sys.stderr, 'Connecting to the database...', dbname
    conn = database.openDb(database.ConnectInfo(dbname))

    if opts.verbose:
        print >> sys.stderr, 'Login in as', username, "in project", project
        
    admin.login(conn, username, project)

    if opts.verbose:
        print >> sys.stderr, 'Loading all plugins...'

    pimgr.loadAllPlugins()

    pmgr = admin.ProjectMgr(conn)
    lmgr = database.LibraryMgr(conn)

    footages = []
    startComposition = sys.maxint
    endComposition = -sys.maxint-1

    for sequence in sequences:

        if opts.verbose:
            print >> sys.stderr, 'Processing layer:', sequence.compName

        # generate list of files
        if not os.path.isabs(sequence.pattern):
            pattern = os.path.join(baseFolder, sequence.pattern)
        else:
            pattern = sequence.pattern

        filelist = genFileList(pattern,
                               sequence.firstFrame,
                               sequence.lastFrame)

        if opts.verbose:
            print >> sys.stderr, 'File list:'
            for fn in filelist:
                print >> sys.stderr, '   ', fn

        # Grab the first part of the filename to make the dbpath
        dbpath = posixpath.join(
            pmgr.getFolder(), project, footageRoot, sequence.compName)       

        if opts.verbose:
            print >> sys.stderr, 'Footage path:', dbpath

        if lmgr.exists(dbpath) and opts.recreate:
            if opts.verbose:
                print '%s exists, deleting it...' % dbpath
            lmgr.remove(dbpath)

        if not lmgr.exists(dbpath):
            dn = posixpath.dirname(dbpath)
            if not lmgr.exists(dn):
                lmgr.makeFolders(dn)

            if opts.verbose:
                print >> sys.stderr, 'Creating layer comp:', dbpath

            # create the composition
            footage = graph.createComp(conn, dbpath, False)

            footage.setWidth(sequence.width)
            footage.setHeight(sequence.height)
            footage.setComponents(sequence.nbComponents)
            footage.setDepth(sequence.depth)
            footage.setPixelRatio((sequence.pixelRatioNum,
                                   sequence.pixelRatioDenom))
            footage.setRate(rate)
            footage.setStart(sequence.timeStart)
            footage.setEnd(sequence.timeEnd)
        else:
            if opts.verbose:
                print >> sys.stderr, 'Updating layer comp:', dbpath

            # get the composition
            footage = graph.getComp(conn, dbpath, True, False)

        # import the files
        desc = '"%s" Imported from %s on %s' % \
               (sequence.renderer, sequence.pattern, datetime.datetime.now())
        result_name = "%s_<Date>_<Time>" % sequence.compName
        if opts.verbose:
            print >> sys.stderr, 'Rendered Result Name:', result_name
            print >> sys.stderr, 'Description:', desc

        rr = graph.referencePublishNoFiles(
            footage, filelist, result_name, 
            sequence.timeStart, sequence.timeEnd, rate,
            sequence.width, sequence.height,
            sequence.depth, sequence.nbComponents,
            (sequence.pixelRatioNum, sequence.pixelRatioDenom), 
            desc)

        footage.setCurrentResult(rr)

        if startComposition > sequence.timeStart:
            startComposition = sequence.timeStart
        if endComposition < sequence.timeEnd:
            endComposition = sequence.timeEnd

        footages.append((sequence, footage))

    # Now create the comp with all the footage and comp nodes
    dbpath = posixpath.join(pmgr.getFolder(), project, compName)

    if lmgr.exists(dbpath) and opts.recreate:
        if opts.verbose:
            print '%s exists, deleting it...' % dbpath
        lmgr.remove(dbpath)

    if not lmgr.exists(dbpath):

        if opts.verbose:
            print >> sys.stderr, 'Creating master comp:', dbpath

        # create the composition
        comp = graph.createComp(conn, dbpath, False)

        if opts.verbose:
            print >> sys.stderr, "Setting rate:", rate
            print >> sys.stderr, "Setting start/end: %d/%d" % \
                  (startComposition, endComposition)
            
        comp.setRate(rate)
        comp.setStart(startComposition)
        comp.setEnd(endComposition )

        linkNodes = []
        for footage in footages:
            linkNodes.append(comp.getGroup().addLinkNode(footage[1]))
        toolReg = graph.ToolRegistry(conn)
        back = linkNodes[0]
        for i in range(len(footages) - 1):
            front = linkNodes[i+1]

            # Creates a blend and comp node with a unique name.
            blendNode = comp.getGroup().createNode(toolReg.getTool('Blend & Comp'))
            if (i == 0):
                blendNode.setName("Blend & Comp")
            else:
                blendNode.setName("Blend & Comp (%i) " % (i,))

            # Set the blend mode
            blendNode['Blend'].setValue(blendMap[footages[i+1][0].blendMode])

            # Connect
            back['Output Image'] >> blendNode['Back']
            front['Output Image'] >> blendNode['Front']
            
            back = blendNode
        
        # Connect input
        back['Output Image'] >> comp.getGroup().getNodes('Output')[0]['Input Image']
 


#------------------------------------------------------------------------------
#
def readImsqfile(imsqfn, opts):
    """
    Reads the IMSQ file and return a dictionary with all the parameters that we
    need extracted from it.
      try to read the IMSQ file and find the frame start and end.

     The 2.0 format of the IMSQ file looks something like this::

    <?xml version="1.0" encoding="utf-8" ?>
    <Sequences version='2.0'>
       <Header>
          <Creator>Maya</Creator>
          <DestFolder>Maya</DestFolder>
          <Project>project</Project>
          <Server>serverName</Server>
          <User>userName</User>
          <Schema>myDatabaseName</Schema>
          <Port>1521</Port>
       </Header>
       <ImageSequence>
          <UniqueName>layer1</UniqueName>
          <LayerName>layer1</LayerName>
          <CameraName>camera3</CameraName>
          <Pattern><![CDATA[layer1/untitled.iff]]></Pattern>
          <Renderer>mayaSoftware</Renderer>
          <FirstFrame>0</FirstFrame>
          <LastFrame>0</LastFrame>
          <FrameIncrement>1</FrameIncrement>
          <TimeStart>0</TimeStart>
          <TimeEnd>0</TimeEnd>
          <TimeStep>1</TimeStep>
          <Width>640</Width>
          <Height>480</Height>
          <PixelAspect>1.000</PixelAspect>
          <TransferMode>Normal</TransferMode>
          <CompositingOrder>1</CompositingOrder>
       </ImageSequence>
       ...
    </Sequences>
    """

    def _getSingleContents(n, name):
        tag = n.getElementsByTagName(name)
        if len(tag) == 0:
            raise RuntimeError("Missing '%s' in '%s' section." %
                               (name, n.tagName))
        if len(tag) != 1 :
            raise RuntimeError("Too many '%s' sections." % name)

        # Un-unicodify the string.
        if len(tag[0].childNodes) == 0 :
            return ""
        return str(tag[0].childNodes[0].wholeText)

    class _Header:
        def __init__(self, xmlNode):
            self.schema = _getSingleContents(xmlNode, 'Schema')
            self.server = _getSingleContents(xmlNode, 'Server')
            self.port = int(_getSingleContents(xmlNode, 'Port'))
            self.username = _getSingleContents(xmlNode, 'User')
            self.project = _getSingleContents(xmlNode, 'Project')
            self.importRoot = _getSingleContents(xmlNode, 'DestFolder')
            self.rate = (int(_getSingleContents(xmlNode, 'Rate_Num')),
                         int(_getSingleContents(xmlNode, 'Rate_Denum')))
            self.compName = _getSingleContents(xmlNode, 'CompositionName')
            self.imageDir = _getSingleContents(xmlNode, 'ImageDir')

        def __str__(self):
            return 'Heading'

    class _ImageSequence:
        def __init__(self, xmlNode):
            self.compName = _getSingleContents(xmlNode, 'UniqueName')

            self.pattern = _getSingleContents(xmlNode, 'Pattern')

            # FIXME: Frames are indexed by ints and time coded by ints inside Toxik.
            # The exact time is derived from the rate. Maya's frame indexes are floats.
            self.firstFrame = int(math.ceil(float(_getSingleContents(xmlNode, 'FirstFrame'))))
            self.lastFrame = int(math.floor(float(_getSingleContents(xmlNode, 'LastFrame'))))
            if self.lastFrame < self.firstFrame:
                raise RuntimeError("Invalid first/last frames: %d/%d" %\
                                   (self.firstFrame, self.lastFrame))
            # Toxik's end frame index is exclusive while it is inclusive in Maya.
            self.lastFrame += 1
            if (self.lastFrame - self.firstFrame > 1 and
                not searchPattern.search(self.pattern)):
                raise RuntimeError(
                    "File pattern requires a ### section for the file number.")

            self.frameIncrement = int(_getSingleContents(xmlNode, 'FrameIncrement'))

            self.timeStart = int(_getSingleContents(xmlNode, 'TimeStart'))            
            self.timeEnd = int(_getSingleContents(xmlNode, 'TimeEnd'))
            if self.timeEnd < self.timeStart:
                raise RuntimeError("Invalid start/end times: %d/%d" %\
                                   (self.timeStart, self.timeEnd))
            # Toxik's end time is exclusive while it is inclusive in Maya.
            self.timeEnd += 1

            self.timeStep = float(_getSingleContents(xmlNode, 'TimeStep'))
            self.width = int(_getSingleContents(xmlNode, 'Width'))
            self.height = int(_getSingleContents(xmlNode, 'Height'))
            self.pixelRatioNum = int(_getSingleContents(xmlNode, 'PixelAspectNum'))
            self.pixelRatioDenom = int(_getSingleContents(xmlNode, 'PixelAspectDenum'))
            self.blendMode = _getSingleContents(xmlNode, 'TransferMode')
            self.compositingOrder = int(_getSingleContents(xmlNode, 'CompositingOrder'))
            self.renderer = _getSingleContents(xmlNode, 'Renderer')

            self.components = _getSingleContents(xmlNode, 'Components')
            if self.components == 'RGBZ':
                self.components = 'RGB'
                print >> sys.stderr, 'Warning: RGBZ images are not yet supported, ignoring Z channel'
            if self.components == 'RGBAZ':
                self.components = 'RGBA'
                print >> sys.stderr, 'Warning: RGBAZ images are not yet supported, ignoring Z channel'			
            self.nbComponents = len(self.components)
            
            pixelDepth = _getSingleContents(xmlNode, 'PixelDepth')
            if pixelDepth == "16 bits int":
                self.depth = 16
            else:
                self.depth = 8

            if not searchPattern.search(self.pattern):
                print >> sys.stderr, "Warning: Non animated sequence "\
                      "(single frame) detected on layer \"%s\". "\
                      "Overwriting start/end with 1/2." % self.compName
                self.timeStart = 1
                self.timeEnd = 2
                self.firstFrame = 1
                self.lastFrame = 2

        def __str__(self):
            return str(self.compositingOrder)

    parsedList = []

    try:
        # Note: since the parsing is really not that involved at all, we use the
        # lame minidom XML parser to avoid the dependency to better, but
        # site-specific library installations.

        dom = domparse(open(imsqfn))
        if len(dom.childNodes) != 1:
            raise RuntimeError("Empty IMSQ file.")
        root = dom.childNodes[0]

        # Make sure we have a sequence document of the right version.
        if root.tagName != 'Sequences':
            raise RuntimeError("No Sequence tag found.")
        if root.getAttribute('version') != '2.0':
            raise RuntimeError("This script only supports IMSQ 2.0.")
       
        # There shouldn't be more than one header
        header = root.getElementsByTagName('Header')
        if len(header) > 1:
            raise RuntimeError("More than one header found.")

        parsedList.append(_Header(header[0]))

        nodeList = root.getElementsByTagName("ImageSequence")

        for node in nodeList:
            parsedList.append(_ImageSequence(node))

        def _itemComparator(x, y):
            if isinstance(x, _Header):
                return -1
            elif isinstance(y, _Header):
                return 1
            if x.compositingOrder == y.compositingOrder:
                raise RuntimeError(
                    "No support for multiple layers at same compositing level.")
            return x.compositingOrder - y.compositingOrder

        parsedList.sort(_itemComparator)

        if opts.verbose:
            print >> sys.stderr, 'Layers found:'
            for layer in parsedList[1:]:
                print >> sys.stderr, '   ', layer.compName

        dom.unlink()

    except IOError, e:
        raise RuntimeError("Error reading IMSQ file '%s': %s" %
                           (imsqfn, str(e)))

    except Exception, e:
        raise RuntimeError("Error parsing IMSQ file \"%s\": %s" % (imsqfn, e))

    return parsedList

#------------------------------------------------------------------------------
#
def parseOptions():
    """
    Parse command-line options and return an opts, args tuple with validated
    options.
    """
    import optparse
    parser = optparse.OptionParser(__doc__.strip())

    parser.add_option('-v', '--verbose', action='store_true',
                      help="Increase verbosity.")

    parser.add_option('-t', '--test', action='store_true',
                      help="Test only, do not execute the Toxik commands.")

    parser.add_option('-p', '--toxikPath', action='store',
                      help="Location where Toxik is installed.")

    parser.add_option('-r', '--recreate', action='store_true',
                      help="Recreate entries already present in the database.")
    
    parser.add_option('-T', '--tempFolder', action='store',
                      help="Temporary folder used for media cache. Needs to exist.")

    opts, args = parser.parse_args()

    if len(args) != 1:
        raise RuntimeError("You must specify an IMSQ filename.")

    return args[0], opts

#------------------------------------------------------------------------------
# STUBS FOR DEBUGGING
def toxik_Stub():
    class ProjectMgrStub:
        def getFolder(self):                      return ""
    class LibraryMgrStub:
        def exists(self, dbname):                 return True
    class referenceStub:
        def getNodes(self):                       return {}
        def __getitem__(self, val):               return self
        def setValue(self, val):                  return self
    class toxikStub:
        def ConnectInfo(self, dbname):            return ""
        def openDb(self, name):                   return ""
        def login(self, conn, user, project):     return ""
        def ProjectMgr(self, conn):               return ProjectMgrStub()
        def LibraryMgr(self, conn):               return LibraryMgrStub()
        def loadAllPlugins(self):                 return ""
        def getComp(self, conn, db, pub, f):      return ""
        def referencePublish(self, f, file, res): return referenceStub()
        
    return toxikStub();

#------------------------------------------------------------------------------
#
def main():

    imsqFile, opts = parseOptions()

    vals = readImsqfile(imsqFile, opts)
    baseFolder = vals[0].imageDir
    if baseFolder == "":
        (baseFolder, filename) = os.path.split(imsqFile)
        if baseFolder == None:
            baseFolder = ""

    createComps("%s@%s:%s" % (vals[0].schema, vals[0].server, vals[0].port),
                vals[0].project, vals[0].username,
                vals[0].importRoot, vals[0].compName, vals[0].rate,
                baseFolder, vals[1:], opts)

#------------------------------------------------------------------------------
#
if __name__ == '__main__':
    try:
        main()
        print >> sys.stderr, 'Toxik: Finished succesfully.'
    except:
        import traceback; traceback.print_exc(file=sys.stderr)
        print >> sys.stderr, 'Toxik: Import failed.'