from pygments.lexer import ExtendedRegexLexer
from pygments.token import *

class CustomLexer(ExtendedRegexLexer):
    name = 'Dig'
    aliases = ['dig']
    filenames = ['*.zone']

    def parse_record(lexer, match, ctx):
        yield match.start('zone'), Name.Function, match.group('zone')
        if match.group('ttl'):
            yield match.start('ttl'), Number.Integer, match.group('ttl')
        yield match.start('cls'), Name.Class, match.group('cls')
        yield match.start('typ'), Name.Tag, match.group('typ')
        yield match.start('value'), String.Escape, match.group('value')
        if match.group('extval'):
            yield match.start('extval'), String.Double, match.group('extval')
        yield match.start('comment'), Comment.Special, match.group('comment')
        ctx.pos = match.end()
    
    tokens = {
        'root': [
            (r'; <<>>.*\n', Comment.Hashbang),  # dig header
            (r';;.*', Generic.Subheading),  # comment section headers
            (r';.*', Comment.Single),  # comment section headers
            (r'(?P<zone>(.*)?\.\s+)(?P<ttl>[0-9]+\s+)?(?P<cls>[A-Z0-9]+\s+)(?P<typ>[A-Z0-9]+\s+)(?P<value>.+?)(?P<extval>\s+?\(\s*?(\n.*?(;[^\n]*)?)+?\))?(?P<comment>(\s+?;.*)?\n)', parse_record),
        ]
    }