
# Copyright 2015 Virantha Ekanayake All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Certain portions of this file (namely the nsprefixes) are borrowed from
# the python-docx project and are provided under the following copyright:
# Copyright (c) 2009-2010 Mike MacCana

#Permission is hereby granted, free of charge, to any person
#obtaining a copy of this software and associated documentation
#files (the "Software"), to deal in the Software without
#restriction, including without limitation the rights to use,
#copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the
#Software is furnished to do so, subject to the following
#conditions:

#The above copyright notice and this permission notice shall be
#included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
#HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
#WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
#OTHER DEALINGS IN THE SOFTWARE.


import zipfile, re
import os,shutil
from plugin import Plugin
from lxml import etree
import logging
import itertools
import copy
from collections import OrderedDict
import tempfile

from plugin import Plugin

class WordResume(Plugin):

    template_file_extension = 'docx'
    nsprefixes = {
    'mo': 'http://schemas.microsoft.com/office/mac/office/2008/main',
    'o':  'urn:schemas-microsoft-com:office:office',
    've': 'http://schemas.openxmlformats.org/markup-compatibility/2006',
    # Text Content
    'w':   'http://schemas.openxmlformats.org/wordprocessingml/2006/main',
    'w10': 'urn:schemas-microsoft-com:office:word',
    'wne': 'http://schemas.microsoft.com/office/word/2006/wordml',
    # Drawing
    'a':   'http://schemas.openxmlformats.org/drawingml/2006/main',
    'm':   'http://schemas.openxmlformats.org/officeDocument/2006/math',
    'mv':  'urn:schemas-microsoft-com:mac:vml',
    'pic': 'http://schemas.openxmlformats.org/drawingml/2006/picture',
    'v':   'urn:schemas-microsoft-com:vml',
    'wp':  'http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing',
    # Properties (core and extended)
    'cp':  'http://schemas.openxmlformats.org/package/2006/metadata/core-properties',
    'dc':  'http://purl.org/dc/elements/1.1/',
    'ep':  'http://schemas.openxmlformats.org/officeDocument/2006/extended-properties',
    'xsi': 'http://www.w3.org/2001/XMLSchema-instance',
    # Content Types
    'ct':  'http://schemas.openxmlformats.org/package/2006/content-types',
    # Package Relationships
    'r':   'http://schemas.openxmlformats.org/officeDocument/2006/relationships',
    'pr':  'http://schemas.openxmlformats.org/package/2006/relationships',
    # Dublin Core document properties
    'dcmitype': 'http://purl.org/dc/dcmitype/',
    'dcterms':  'http://purl.org/dc/terms/'}

    char_loop_start = '<'
    char_loop_end   = '>'
    char_tag_start  = '['
    char_tag_end    = ']'

    def __init__ (self, template_file, resume_data, skip):
        self.skip = skip
        self.resume_data = resume_data
        self.template_filename = template_file

        # Open the docx file and read in the xml tree
        self.template = open (self.template_filename)
        self.doc_etree = self.read_contents(self.template)

    def render(self, output_filename):
        """
            Needs read_contents to have been called
        """
        self.substitute_data()
        self.write_and_close_docx(self.doc_etree, output_filename)

    def read_contents(self, filename):
        """
            Unzip and read the contents of the word/document.xml file

            :rtype: the contents of the file as an ElementTree
        """
        self.zipfile = zipfile.ZipFile(filename)
        xml_content = self.zipfile.read('word/document.xml')
        return etree.fromstring(xml_content)

    def _check_element_is(self, element, type_char):
        return element.tag == '{%s}%s' % (self.nsprefixes['w'],type_char)
    def _assert_element_is(self, element, type_char):
        assert self._check_element_is(element, type_char)

    def _get_all_text_in_node(self, node):
        all_txt = []
        for txt in node.itertext(tag=etree.Element):
            all_txt.append(txt)
        return ''.join(all_txt)

    def _get_parent_paragraph(self, text_node):
        self._assert_element_is(text_node, 't')
        run = text_node.getparent()
        self._assert_element_is(run, 'r')
        paragraph = run.getparent()
        self._assert_element_is(paragraph, 'p')
        return paragraph


    def _extract_loop(self, paragraph):
        """
            Find and copy the xml tree fragment that represents the '<' and '>' delimited
            loop.
        """
        inside_loop = False
        loop_tree = etree.Element("root")  # Create a new tree root
        prev_paragraph = None
        list_of_all_loop_nodes = []

        # Keep looking for the first (and only) loop after this tag
        for node, text, node_index in self._itersiblingtext(paragraph):
            if '<' in text:
                assert not inside_loop
                inside_loop = True
                loop_start_node = self._get_parent_paragraph(node)
                self._assert_element_is(loop_start_node, 'p')

            if inside_loop:
                # This text node is enclosed by a run, which is in turn enclosed by the paragraph
                # This paragraph is what we want to extract
                current_paragraph = self._get_parent_paragraph(node)
                self._assert_element_is(current_paragraph, 'p')
                if prev_paragraph != current_paragraph:
                    loop_tree.insert(node_index, copy.deepcopy(current_paragraph))
                    prev_paragraph = current_paragraph
                    list_of_all_loop_nodes.append(current_paragraph)

            if '>' in text:
                assert inside_loop
                # Done with finding loop, so exit this iterator
                break
        logging.debug("Found a loop spanning %d paragraphs, %d" % (node_index, len(loop_tree)))
        return loop_start_node, loop_tree, list_of_all_loop_nodes

    def _make_loop_instance(self, loop_tree, section_data, all_tagnames_in_section):

        # Copy the loop_tree to make an instance we can fill in with the results from 
        # the section names
        loop_instance = copy.deepcopy(loop_tree)

        for node, text, node_i in self._itersiblingtext(loop_instance):
            if '<' in text: node.text = node.text.replace('<','')
            if '>' in text: node.text = node.text.replace('>','')

            # For every node with text in it, don't even try to find a tag name
            # Instead, brute force it by trying to sub every possible tag in every node
            if node.text:
                for key in all_tagnames_in_section:
                    logging.debug("Looking for key %s" % key)
                    instance_text = section_data.get(key, '')  # Replace any key that isn't specified for this instance with a blank
                    node.text = node.text.replace('['+key+']', str(instance_text))
        return loop_instance

    def _fill_in_section(self, section_etree, section_list):
        """
            section_list is a list of dicts containing the contents for the current section. For example:
                
                [
                 { key1:val1, key2:val2 },
                 { key1:val1, key2:val2, key3:val3},
                ]

            Assumptions:
                - All loop starts are in their own paragraphs
                - loops can span multiple paragraphs
                - Any content after the loop must be in different paragraph
                - No tables anywhere!
        """
        # Get the parent paragraph 
        paragraph = self._get_parent_paragraph(section_etree)

        # Build up a copy of the loop sub-tree in loop_tree. We will instance this as many times as needed.
        # Also, keep track of all the elements belonging to the loop body in elements_to_delete
        # so that we can delete the original loop definition once we are done instancing the loop.
        loop_start_node, loop_tree, elements_to_delete = self._extract_loop(paragraph)
        
        # Max possible set of subtags
        all_tag_names_in_section = self._get_all_keys_in_list_of_dicts(section_list)

        # Iterate and create the loop with the section data
        loop_instance = []
        for section_data in section_list:
            loop_instance.append(self._make_loop_instance(loop_tree, section_data, all_tag_names_in_section))

        # Now, add the loop_instance to the body
        body = loop_start_node.getparent()
        index_to_insert_at = body.index(loop_start_node)+1
        logging.debug("Inserting at index %d, %d instances" % (index_to_insert_at, len(loop_instance)))
        for inst in loop_instance:
            for child in inst.getchildren():
                body.insert(index_to_insert_at, child)
                index_to_insert_at += 1

        # Delete the loop template
        for e in elements_to_delete:
            body.remove(e)

        # Remove a section heading name if it's preceded with a bang (for example, a contact section, where
        # we don't need to title the section
        if '[!' in section_etree.text:
            paragraph.getparent().remove(paragraph)

    def _find_sections(self, my_etree, tags_to_find, char_to_stop_on=None):
        """
            Build a dict of all the top-level tags in the document and the corresponding
            text node.

            Also replace the top level tag header with the section name

            :rtype: Dict of tag name -> xml node
            
        """

        tags = {}
        logging.debug("Looking for tags: %s" % (','.join(tags_to_find)))
        mTag = r"""\[\!?(?P<tag>[\s\w\_\|]+)\]"""
        # lowercase all the tags_to_find, just in case the user upper-cased some of them
        tags_to_find = [x.lower() for x in tags_to_find]

        for node,text in self._itertext(my_etree):
            tag_text = re.findall(mTag, text) 
            if tag_text:
                logging.debug("Found grps %s" % (','.join(tag_text)))
            for tag in tag_text:
                tag_lower = tag.lower()
                tag_lower = tag_lower.split('|')[0]
                if tag_lower in tags_to_find:
                    tags[tag_lower] = node
                    # Replace the brackets with nothing
                    if '[!' in node.text:
                        #node.text = node.text.replace('[!'+tag+']', 'lsdjflkd')
                        body = node.getparent().getparent().getparent()
                        #body.remove(node.getparent().getparent())
                    else:
                        node.text = node.text.replace('['+tag+']', tag)
                        if '|' in node.text:
                            # We have alternate text so just use that
                            # FIXME:  This is wrong, we need to sub out the text, and not replace the whole node!
                            node.text = node.text.split('|')[1]

        return tags

    def _get_all_keys_in_list_of_dicts(self, mylist):
        mykeys = set()
        for e in mylist:
            for k in e.keys():
                mykeys.add(k)
        return list(mykeys)

    def substitute_data(self):

        body = self.doc_etree.xpath('/w:document/w:body', namespaces=self.nsprefixes)[0]
        self.collapse_tags(body)

        if self.skip:
            return

        # Get a list of all the top-level tags
        sections = self._find_sections(self.doc_etree, self.resume_data.keys())
        for section_name, section_node in sections.items():
            logging.debug("Subtag search for %s" % section_name)
            self._fill_in_section(section_node, self.resume_data[section_name])
            logging.debug("Finished find_subtags_in_loop")
        return

    def _itertext(self, my_etree):
        """Iterator to go through xml tree's text nodes"""
        for node in my_etree.iter(tag=etree.Element):
            if self._check_element_is(node, 't'):
                yield (node, node.text)

    def _itersiblingtext(self, my_etree):
        """Iterator to go through xml tree sibling text nodes"""
        # First, check if there are siblings, if not, we need to just go through the current element
        if my_etree.getnext() is None:
            for node in my_etree.iter(tag=etree.Element):
                if self._check_element_is(node, 't'):
                    yield (node, node.text, 0)
        else:
            for i, sib in enumerate(my_etree.itersiblings(tag=etree.Element)):
                for node in sib.iter(tag=etree.Element):
                    if self._check_element_is(node, 't'):
                        yield (node, node.text, i)


    def collapse_tags(self, my_etree):
        """
            Find the special tags and collapse all the text in them to one single paragraph. This
            method works in-place, modifying the passed in etree.
            This works because we know there shouldn't be any formatting inside the tag name.


            :rtype: Nothing
        """
        chars = []
        is_tag_start = False      # True if inside tag
        tag_start_node = None     # Pointer to current node. 
        tag_start_char = '['
        tag_end_char = ']'

        # For every node with text
        for node,text in self._itertext(my_etree):
            # Go through each node's text character by character
            for i,c in enumerate(text):
                if c == tag_start_char:  # Tag is starting!
                    assert not is_tag_start  # Better not already be inside a tag!
                    is_tag_start = True 
                    tag_start_node = node 
                    chars = []
                elif c == tag_end_char:  # Tag is ending
                    assert is_tag_start  # Better have seen a tag start!
                    is_tag_start = False
                    # If tag_start_node is the same as current node, then we don't need to do anything
                    # But otherwise:
                    if node != tag_start_node:
                        # Tag started in different node, so move all the chars we've encountered since then
                        # to the tag_start_node
                        chars.append(c)
                        tag_start_node.text += ''.join(chars)
                        node.text = text[i+1:]   # Remove characters from this node
                else:
                    # Normal text character
                    if is_tag_start and node != tag_start_node:
                        # Need to save these chars to append to text in openbrac_node
                        chars.append(c)

            # If we're here, that means we've consumed all the text in the current node.
            # Check if this node was part of a tag, yet did not start the tag
            if is_tag_start and node!= tag_start_node:
                # Need to remove this text completely as we've saved all of it inside chars for moving
                # into the start_node
                node.text = ""


    def write_and_close_docx (self, xml_content, output_filename):
        """ Create a temp directory, expand the original docx zip.
            Write the modified xml to word/document.xml
            Zip it up as the new docx
        """

        tmp_dir = tempfile.mkdtemp()

        self.zipfile.extractall(tmp_dir)

        with open(os.path.join(tmp_dir,'word/document.xml'), 'w') as f:
            xmlstr = etree.tostring (xml_content, pretty_print=True)
            f.write(xmlstr)

        # Get a list of all the files in the original docx zipfile
        filenames = self.zipfile.namelist()

        # Now, create the new zip file and add all the filex into the archive
        zip_copy_filename = output_filename
        with zipfile.ZipFile(zip_copy_filename, "w") as docx:
            for filename in filenames:
                docx.write(os.path.join(tmp_dir,filename), filename)

        # Clean up the temp dir
        shutil.rmtree(tmp_dir)


if __name__ == '__main__':
    script = WordAPI('test/blah.docx', {}, False)
    print (etree.tostring(script.doc_etree, pretty_print=True))
    script.collapse_tags(script.doc_etree)
    print (etree.tostring(script.doc_etree, pretty_print=True))


