#!/usr/bin/python

import yaml, codecs, sys, os.path, optparse

class Style:

    def __init__(self, header=None, footer=None,

            tokens=None, events=None, replaces=None):

        self.header = header

        self.footer = footer

        self.replaces = replaces

        self.substitutions = {}

        for domain, Class in [(tokens, 'Token'), (events, 'Event')]:

            if not domain:

                continue

            for key in domain:

                name = ''.join([part.capitalize() for part in key.split('-')])

                cls = getattr(yaml, '%s%s' % (name, Class))

                value = domain[key]

                if not value:

                    continue

                start = value.get('start')

                end = value.get('end')

                if start:

                    self.substitutions[cls, -1] = start

                if end:

                    self.substitutions[cls, +1] = end

    def __setstate__(self, state):

        self.__init__(**state)

yaml.add_path_resolver(u'tag:yaml.org,2002:python/object:__main__.Style',

        [None], dict)

yaml.add_path_resolver(u'tag:yaml.org,2002:pairs',

        [None, u'replaces'], list)

class YAMLHighlight:

    def __init__(self, options):

        config = yaml.load(file(options.config, 'rb').read())

        self.style = config[options.style]

        if options.input:

            self.input = file(options.input, 'rb')

        else:

            self.input = sys.stdin

        if options.output:

            self.output = file(options.output, 'wb')

        else:

            self.output = sys.stdout

    def highlight(self):

        input = self.input.read()

        if input.startswith(codecs.BOM_UTF16_LE):

            input = unicode(input, 'utf-16-le')

        elif input.startswith(codecs.BOM_UTF16_BE):

            input = unicode(input, 'utf-16-be')

        else:

            input = unicode(input, 'utf-8')

        substitutions = self.style.substitutions

        tokens = yaml.scan(input)

        events = yaml.parse(input)

        markers = []

        number = 0

        for token in tokens:

            number += 1

            if token.start_mark.index != token.end_mark.index:

                cls = token.__class__

                if (cls, -1) in substitutions:

                    markers.append([token.start_mark.index, +2, number, substitutions[cls, -1]])

                if (cls, +1) in substitutions:

                    markers.append([token.end_mark.index, -2, number, substitutions[cls, +1]])

        number = 0

        for event in events:

            number += 1

            cls = event.__class__

            if (cls, -1) in substitutions:

                markers.append([event.start_mark.index, +1, number, substitutions[cls, -1]])

            if (cls, +1) in substitutions:

                markers.append([event.end_mark.index, -1, number, substitutions[cls, +1]])

        markers.sort()

        markers.reverse()

        chunks = []

        position = len(input)

        for index, weight1, weight2, substitution in markers:

            if index < position:

                chunk = input[index:position]

                for substring, replacement in self.style.replaces:

                    chunk = chunk.replace(substring, replacement)

                chunks.append(chunk)

                position = index

            chunks.append(substitution)

        chunks.reverse()

        result = u''.join(chunks)

        if self.style.header:

            self.output.write(self.style.header)

        self.output.write(result.encode('utf-8'))

        if self.style.footer:

            self.output.write(self.style.footer)

if __name__ == '__main__':

    parser = optparse.OptionParser()

    parser.add_option('-s', '--style', dest='style', default='ascii',

            help="specify the highlighting style", metavar='STYLE')

    parser.add_option('-c', '--config', dest='config',

            default=os.path.join(os.path.dirname(sys.argv[0]), 'yaml_hl.cfg'),

            help="set an alternative configuration file", metavar='CONFIG')

    parser.add_option('-i', '--input', dest='input', default=None,

            help="set the input file (default: stdin)", metavar='FILE')

    parser.add_option('-o', '--output', dest='output', default=None,

            help="set the output file (default: stdout)", metavar='FILE')

    (options, args) = parser.parse_args()

    hl = YAMLHighlight(options)

    hl.highlight()