# -*- coding: utf-8 -*-
# Copyright (c) 2012-2015, Anima Istanbul
#
# This module is part of anima-tools and is released under the BSD 2
# License: http://www.opensource.org/licenses/BSD-2-Clause

import os
import subprocess
import tempfile
import platform
from contextlib import contextmanager

import pymel.core as pm
from pymel.core.uitypes import CheckBox, TextField
from pymel.core.general import Attribute
from pymel.core.system import FileReference

from anima.edit import Sequence, Media, Video, Track, Clip, File, Rate
from anima.extension import extends
from anima.repr import Representation


default_handle_count = 0


class MayaExtension(object):
    """Extension to PyMel classes
    """

    @extends(Attribute)
    @property
    def next_available(self):
        """returns the next available attr in a multi attr

        :return: The available index as an attribute
        """
        try:
            indices = self.getArrayIndices()
        except TypeError:
            return self

        available_index = 0

        try:
            for i in xrange(max(indices) + 2):
                if not self[i].connections():
                    available_index = i
                    break
        except ValueError:
            available_index = 0

        return self[available_index]

    @extends(CheckBox)
    def value(self, value=None):
        """returns or set the check box value
        """
        from pymel.core import checkBox
        if value is not None:
            # set the value
            checkBox(self, e=1, v=value)
        else:
            # get the value
            return checkBox(self, q=1, v=1)

    @extends(TextField)
    def text(self, value=None):
        """returns or sets the text field value
        """
        from pymel.core import textField
        if value is not None:
            # set the value
            textField(self, e=1, tx=value)
        else:
            # get the value
            return textField(self, q=1, tx=1)


class FileReferenceExtension(object):
    """Extensions to Maya FileReference node.

    Manages the Referenced different representations in the current scene.

    This class helps converting references to other representations and vice
    versa.
    """

    def path(self):
        """dummy method to let the IDEs find FileReference.path and not
        complain about it.
        """
        return ''

    def replaceWith(self, path):
        """dummy method to let the IDEs find FileReference.replaceWith and not
        complain about it.
        """
        return None

    @extends(FileReference)
    def to_repr(self, repr_name):
        """Replaces the current reference with the representation with the
        given repr_name.

        :param str repr_name: The desired repr name
        :return:
        """
        rep_v = self.find_repr(repr_name)
        from stalker import Repository
        if rep_v is not None and rep_v != self.version:
            self.replaceWith(
                Repository.to_os_independent_path(rep_v.absolute_full_path)
            )

    @extends(FileReference)
    def find_repr(self, repr_name):
        """Finds the representation with the given repr_name.

        :param str repr_name: The desired repr name
        :return: :class:`.Version`
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return

        rep = Representation(version=v)
        rep_v = rep.find(repr_name)

        return rep_v

    @extends(FileReference)
    def list_all_repr(self):
        """Returns a list of strings representing all the representation names
        of this FileReference

        :return: list of str
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return []

        rep = Representation(version=v)
        return rep.list_all()

    @extends(FileReference)
    def to_base(self):
        """Loads the original version
        """
        self.to_repr(Representation.base_repr_name)

    @extends(FileReference)
    def is_base(self):
        """returns True or False depending to if this is the base
        representation for this reference
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return True

        rep = Representation(version=v)
        return rep.is_base()

    @extends(FileReference)
    def has_repr(self, repr_name):
        """checks if the reference has the given representation

        :param str repr_name: The name of the desired representation
        :return:
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return False

        rep = Representation(version=v)
        return rep.has_repr(repr_name)

    @extends(FileReference)
    def is_repr(self, repr_name):
        """returns True or False depending to if this is the requested repr

        :param str repr_name: The representation name
        :return:
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return False

        rep = Representation(version=v)
        return rep.is_repr(repr_name=repr_name)

    @extends(FileReference)
    @property
    def repr(self):
        """the representation name of the related version
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return None

        rep = Representation(version=v)
        return rep.repr

    @extends(FileReference)
    def get_base(self):
        """returns the base version instance
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        v = m.get_version_from_full_path(self.path)

        if v is None:
            return True

        rep = Representation(version=v)
        return rep.find(rep.base_repr_name)

    @extends(FileReference)
    @property
    def version(self):
        """returns the Stalker Version instance related to this reference
        """
        from anima.env.mayaEnv import Maya
        m = Maya()
        return m.get_version_from_full_path(self.path)


class SequenceManagerExtension(object):
    """Extension to the pm.nodetypes.SequenceManager class
    """

    @extends(pm.nodetypes.SequenceManager)
    def get_shot_name_template(self):
        """returns teh shot_name_template_ attribute value, creates the
        attribute if missing
        """
        if not self.hasAttr('shot_name_template'):
            default_template = '<Sequence>_<Shot>_<Version>'
            self.set_shot_name_template(default_template)

        return self.shot_name_template.get()

    @extends(pm.nodetypes.SequenceManager)
    def set_shot_name_template(self, template):
        """sets the shot_name_template attribute value
        """
        if not self.hasAttr('shot_name_template'):
            self.addAttr('shot_name_template', dt='string')

        self.shot_name_template.set(template)

    @extends(pm.nodetypes.SequenceManager)
    def get_version(self):
        """returns the version attribute value, creates the attribute if
        missing
        """
        if not self.hasAttr('version'):
            self.set_version('')
        return self.version.get()

    @extends(pm.nodetypes.SequenceManager)
    def set_version(self, template):
        """sets the version attribute value
        """
        if not self.hasAttr('version'):
            self.addAttr('version', dt='string')

        self.version.set(template)

    @extends(pm.nodetypes.SequenceManager)
    def create_sequence(self, name=None):
        """Creates a new sequence

        :return: pm.nodetypes.Sequence
        """
        sequencer = pm.createNode('sequencer')
        if name:
            sequencer.set_sequence_name(name)

        sequencer.message >> self.sequences.next_available
        return sequencer

    @extends(pm.nodetypes.SequenceManager)
    def from_seq(self, seq):
        """Generates maya shots and sequencers with the given
        :class:`anima.previs.Sequence` instance.

        This method is designed to mainly be used with :method:`.from_xml`
        and :method:`.from_edl` methods. It generates new shots or updates them
        by looking at the given :class:`anima.previs.Sequence` instance.

        :param seq: An :class:`anima.previs.Sequence` instance
        :return:
        """
        # now create or update structure
        # create first
        # generate the shot name template first
        shot_name_template = self.get_shot_name_template()
        # get current sequencer
        seqs = self.sequences.get()
        if seqs:
            # we probably need to update shots
            seq1 = seqs[0]

            # update shots
            shots = seq1.shots.get()

            # collect clips
            all_clips = [clip
                         for track in seq.media.video.tracks
                         for clip in track.clips]

            deleted_shots = []
            used_shots = []

            # for each shot sets the anchor point in a temp variable
            for shot in shots:
                shot.anchor = shot.startFrame.get()

            # find the corresponding shots in seq
            for clip in all_clips:
                #is_deleted = True
                for shot in shots:
                    if clip.id.lower() == shot.full_shot_name.lower():

                        # update with the given clip info
                        # anchor = shot.startFrame.get()
                        handle = shot.handle.get()
                        track = shot.track.get()

                        if shot in used_shots:
                            # this shot has been used once so duplicate it
                            # before doing anything
                            dup_shot = seq1.create_shot(shot.shotName.get())
                            dup_shot.startFrame.set(shot.startFrame.get())
                            dup_shot.endFrame.set(shot.endFrame.get())
                            dup_shot.handle.set(shot.handle.get())
                            dup_shot.sequenceStartFrame.set(
                                shot.sequenceStartFrame.get()
                            )
                            # do not copy sequenceEndFrame
                            # copy camera
                            dup_shot.set_camera(shot.get_camera())
                            dup_shot.anchor = shot.anchor
                            shot = dup_shot

                        start_frame = clip.in_ - handle + shot.anchor
                        end_frame = clip.out - clip.in_ + start_frame - 1

                        # print '-------------------------------'
                        # print 'clip.in     : %s' % clip.in_
                        # print 'clip.out    : %s' % clip.out
                        # print 'clip.start  : %s' % clip.start
                        # print 'clip.end    : %s' % clip.end
                        # print 'start_frame : %s' % start_frame
                        # print 'end_frame   : %s' % end_frame

                        sequence_start = clip.start

                        shot.startFrame.set(start_frame)
                        shot.endFrame.set(end_frame)

                        shot.sequenceStartFrame.set(sequence_start)

                        # set original track
                        shot.track.set(track)

                        used_shots.append(shot)
                        break

            for shot in seq1.shots.get():
                if shot not in used_shots:
                    deleted_shots.append(shot)

            # delete shots
            pm.delete(deleted_shots)

        else:
            # create sequencer
            seq1 = self.create_sequence(seq.name)

            # create shots
            media = seq.media
            for i, track in enumerate(media.video.tracks):
                for clip in track.clips:
                    # clip.id is something like SEQ001_HSNI_010_0010_v046
                    # filter the shot name
                    shot_name = clip.id.split('_')[-2]
                    shot = seq1.create_shot(shot_name)

                    shot.startFrame.set(clip.in_)
                    shot.endFrame.set(clip.out - 1)
                    shot.sequenceStartFrame.set(clip.start)
                    shot.handle.set(0)
                    if clip.file:
                        f = clip.file
                        pathurl = f.pathurl.replace('file://localhost/', '')
                        if ':' not in pathurl:  # not windows, keep '/'
                            pathurl = '/%s' % pathurl
                        shot.output.set(pathurl)

                    shot.track.set(i + 1)

    @extends(pm.nodetypes.SequenceManager)
    def from_xml(self, path):
        """Parses XML file and returns a Sequence instance which reflects the
        whole time line hierarchy.

        :param path: The path of the XML file
        :return: :class:`.Sequence`
        """
        if not isinstance(path, str):
            raise TypeError(
                'path argument in %s.from_xml should be a string, not %s' %
                (self.__class__.__name__, path.__class__.__name__)
            )

        from xml.etree import ElementTree

        try:
            tree = ElementTree.parse(path)
        except IOError:
            raise IOError('Please supply a valid path to an XML file!')

        root = tree.getroot()
        seq = Sequence()
        xml_seq = root.getchildren()[0]

        seq.from_xml(xml_seq)

        self.from_seq(seq)

    @extends(pm.nodetypes.SequenceManager)
    def from_edl(self, path):
        """Parses EDL file and returns a Sequence instance which reflects the
        whole time line hierarchy.

        :param path: The path of the XML file
        :return: :class:`.Sequence`
        """
        if not isinstance(path, str):
            raise TypeError(
                'path argument in %s.from_edl should be a string, not %s' %
                (self.__class__.__name__, path.__class__.__name__)
            )

        from anima.env.mayaEnv import Maya
        m = Maya()
        fps = m.get_fps()

        import edl
        p = edl.Parser(str(fps))
        with open(path) as f:
            l = p.parse(f)

        seq = Sequence()
        seq.from_edl(l)

        self.from_seq(seq)

    @extends(pm.nodetypes.SequenceManager)
    def to_xml(self, indentation=2, pre_indent=0):
        """Generates an FCP compatible XML file at given path.

        :return:
        """
        seq = self.generate_sequence_structure()

        rendered_template = ''
        if seq:
            rendered_template = seq.to_xml(
                indentation=indentation,
                pre_indent=indentation + pre_indent
            )
        return rendered_template

    @extends(pm.nodetypes.SequenceManager)
    def generate_sequence_structure(self):
        """Generates a Sequence structure suitable for XML<->EDL conversion

        :return: Sequence
        """
        import timecode
        from anima.env.mayaEnv import Maya

        m = Maya()
        fps = m.get_fps()

        # export only the first sequence, ignore others
        sequencers = self.sequences.get()
        if len(sequencers) == 0:
            return None

        sequencer = sequencers[0]
        time = pm.PyNode('time1')

        seq = Sequence()
        seq.name = str(sequencer.get_sequence_name())

        seq.rate = Rate(timebase=str(fps), ntsc=False)
        seq.timecode = str(timecode.Timecode(
            framerate=seq.rate.timebase,
            frames=time.timecodeProductionStart.get() + 1
        ))
        seq.duration = sequencer.duration

        media = Media()
        video = Video()
        media.video = video

        for shot in sequencer.shots.get():
            clip = Clip()
            clip.id = str(shot.full_shot_name)
            clip.name = str(shot.full_shot_name)
            clip.duration = shot.duration + 2 * shot.handle.get()
            clip.enabled = True
            clip.start = shot.sequenceStartFrame.get()
            clip.end = shot.sequenceEndFrame.get() + 1
            # clips always start from 0 and includes the shot handle
            clip.in_ = shot.handle.get()  # handle at start
            clip.out = shot.handle.get() + shot.duration  # handle at end
            clip.type = 'Video'  # always video for now

            f = File()
            f.name = os.path.splitext(
                os.path.basename(str(shot.output.get()))
            )[0]

            f.duration = shot.duration + 2 * shot.handle.get()

            f.pathurl = str('file://localhost/%s' % shot.output.get())

            clip.file = f

            track_number = shot.track.get() - 1  # tracks should start from 0
            try:
                track = video.tracks[track_number]
            except IndexError:
                track = Track()
                video.tracks.append(track)

            track.clips.append(clip)
            # set video resolution
            video.width = shot.wResolution.get()
            video.height = shot.hResolution.get()

        seq.media = media
        return seq

    @extends(pm.nodetypes.SequenceManager)
    def to_edl(self):
        """Generates an EDL file out of the edit
        """
        seq = self.generate_sequence_structure()
        l = seq.to_edl()
        return l


class SequencerExtension(object):
    """The sequence instance.

    It is a manager that manages shot data. It is kind of the reflection of the
    Maya Sequencer instance.

    It is able to get Maya editorial XML and convert it to EDL.
    """

    @extends(pm.nodetypes.Sequencer)
    @property
    def manager(self):
        """returns the SequenceManager instance that this sequence is connected
        to
        """
        return pm.ls(self.connections(), type='sequenceManager')[0]

    @extends(pm.nodetypes.Sequencer)
    def get_sequence_name(self):
        """Gets the sequence_name attribute value, creates the attribute if it
        is missing
        """
        if not self.hasAttr('sequence_name'):
            self.set_sequence_name('')
        return self.sequence_name.get()

    @extends(pm.nodetypes.Sequencer)
    def set_sequence_name(self, name):
        """Sets the sequence name for this Sequencer

        :param name: A str holding the desired name of the sequence.
        :return: None
        """
        if not self.hasAttr('sequence_name'):
            self.addAttr('sequence_name', dt='string')
        self.sequence_name.set(name)

    @extends(pm.nodetypes.Sequencer)
    @property
    def all_shots(self):
        """return all the shots connected to this sequencer
        """
        return pm.ls(self.shots.get(), typ='shot')

    # @extends(pm.nodetypes.Sequencer)
    # @all_shots.setter
    # def all_shots(self, shots):
    #     """setter for the all_shots property
    #     """
    #     # remove the current shots first
    #     # then append the new ones
    #     for s in self.all_shots:
    #         t = s // self.shots
    #
    #     for s in shots:
    #         t = s.message >> self.shots.next_available

    @extends(pm.nodetypes.Sequencer)
    def add_shot(self, shot):
        """Adds the given shot to the current sequencer

        :param shot: a pm.nodetypes.Shot instance
        :return: None
        """
        # add the given shot to the list
        if not shot.hasAttr('handle'):
            shot.set_handle(handle=default_handle_count)

        # connect to the sequencer
        # remove it from other sequences
        for attr in shot.message.outputs(p=1):
            if isinstance(attr.node(), pm.nodetypes.Sequencer):
                shot.message // attr

        shot.message >> self.shots.next_available

    @extends(pm.nodetypes.Sequencer)
    def set_shot_handles(self, handle=default_handle_count):
        """Set shot handles

        :param int handle: An integer value for handle
        :return: