#!/usr/bin/python

from __future__ import print_function

import os, shutil, sys
import StringIO

# file creation decorator
def fileMaker(func):
    def inner(*args, **kwargs):
        print(" - Creating file: %s..." % args[1], end='')
        sys.stdout.flush()
        retval = func(*args)
        print("done.")
        sys.stdout.flush()
        return retval
    return inner

# symlink creation decorator
def fileLinker(func):
    def inner(*args, **kwargs):
        print(" - Creating symlink: %s..." % args[1], end='')
        sys.stdout.flush()
        retval = func(*args)
        print("done.")
        sys.stdout.flush()
        return retval
    return inner

# return a copy of original, minues any lines that contain an entry in blacklist
def trimXmlElements(original, blacklist):
    out = StringIO.StringIO()
    for line in original.splitlines():
        if all(x not in line for x in blacklist):
            out.write(line + '\n')
    result = out.getvalue()
    out.close()
    return result

class TestDataGenerator:

    def __init__(self, source, dest):

        # source/destination directories
        self.testDataDir      = source
        self.generatedDataDir = dest

        # generated output files/symlinks & 'maker' functions
        self.outputFiles = {
            'truncated.bam' : self.makeTruncatedBam,
            'chunking_emptyfilters.subreadset.xml'   : self.makeChunkingXml,
            'chunking_missingfilters.subreadset.xml' : self.makeChunkingXml,
        }
        self.outputSymlinks = {
            'aligned.bam'      : self.makeAlignedBamCopy,
            'aligned.bam.bai'  : self.makeAlignedBamCopy,
            'aligned.bam.pbi'  : self.makeAlignedBamCopy,
            'aligned2.bam'     : self.makeAlignedBamCopy,
            'aligned2.bam.bai' : self.makeAlignedBamCopy,
            'aligned2.bam.pbi' : self.makeAlignedBamCopy,
            'm150404_101626_42267_c100807920800000001823174110291514_s1_p0.1.subreads.bam'     : self.makeChunkingSymlink,
            'm150404_101626_42267_c100807920800000001823174110291514_s1_p0.1.subreads.bam.pbi' : self.makeChunkingSymlink,
            'm150404_101626_42267_c100807920800000001823174110291514_s1_p0.2.subreads.bam'     : self.makeChunkingSymlink,
            'm150404_101626_42267_c100807920800000001823174110291514_s1_p0.2.subreads.bam.pbi' : self.makeChunkingSymlink,
            'm150404_101626_42267_c100807920800000001823174110291514_s1_p0.3.subreads.bam'     : self.makeChunkingSymlink,
            'm150404_101626_42267_c100807920800000001823174110291514_s1_p0.3.subreads.bam.pbi' : self.makeChunkingSymlink,
        }

    def editChunkingXml(self, outputFn, removeFiltersNode):
        inputXmlFn  = os.path.join(self.testDataDir,'chunking','chunking.subreadset.xml')
        outputXmlFn = os.path.join(self.generatedDataDir,outputFn)

        blacklist = ['pbds:Filter>', 'pbbase:Properties>', '<pbbase:Property']
        if removeFiltersNode:
            blacklist.append('pbds:Filters>')

        inputXml = ''
        with open(inputXmlFn, 'r') as xml_infile:
            inputXml = xml_infile.read()
        outputXml = trimXmlElements(inputXml, blacklist)
        with open(outputXmlFn, 'w') as xml_outfile:
            xml_outfile.write(outputXml)

    @fileLinker
    def makeAlignedBamCopy(self, outputFn):
        source = os.path.join(self.testDataDir,outputFn)
        dest   = os.path.join(self.generatedDataDir, outputFn)
        os.symlink(source, dest)

    @fileLinker
    def makeChunkingSymlink(self, outputFn):
        source = os.path.join(self.testDataDir,'chunking', outputFn)
        dest   = os.path.join(self.generatedDataDir, outputFn)
        os.symlink(source, dest)

    @fileMaker
    def makeChunkingXml(self, outputFn):
        if outputFn == 'chunking_emptyfilters.subreadset.xml':
            removeFiltersNode = False
        else:
            removeFiltersNode = True
        self.editChunkingXml(outputFn, removeFiltersNode)

    @fileMaker
    def makeTruncatedBam(self, outputFn):
        source = os.path.join(self.testDataDir, 'phi29.bam')
        dest   = os.path.join(self.generatedDataDir, outputFn)
        shutil.copyfile(source, dest)
        with open(dest, 'r+b') as in_file:
            in_file.truncate(200)

    # main entry point
    def generate(self):

        # skip file if it exists
        os.chdir(self.generatedDataDir)
        filenames = self.outputFiles.keys()
        for file in filenames:
            if os.path.exists(file) :
                del self.outputFiles[file]

        # skip symlink if it exists
        symlinks = self.outputSymlinks.keys()
        for link in symlinks:
            if os.path.lexists(link):
                del self.outputSymlinks[link]

        # only print message & run makers, if any files/symlinks to be created
        # else silent success
        if self.outputFiles or self.outputSymlinks:
            print('Generating test data in %s ' % self.generatedDataDir)
            for file, func in self.outputFiles.iteritems():
                func(file)
            for link, func in self.outputSymlinks.iteritems():
                func(link)

# script entry point
if __name__ == '__main__':
    g = TestDataGenerator(sys.argv[1], sys.argv[2])
    g.generate()
