# -*- coding: utf-8 -*-
#
# https://gitlab.com/s3lph/pygments-dig-lexer
#
# SPDX-FileCopyrightText: 2020 s3lph
# SPDX-License-Identifier: MIT

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

REPLY_CODES = {
    'NOERROR': Name.Function,
    'NXDOMAIN': String.Escape,
    'REFUSED': String.Escape,
    'NOTAUTH': String.Escape,
    'NOTZONE': String.Escape
}

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

    def parse_dig_header(lexer, re, ctx):
        yield re.start('c1'), Comment.Hashbang, re.group('c1')
        yield re.start('version'), String.Doc, re.group('version')
        yield re.start('c2'), Comment.Hashbang, re.group('c2')
        yield re.start('query'), Name.Function, re.group('query')
        ctx.pos = re.end()

    def parse_response_header(lexer, re, ctx):
        yield re.start('c1'), Comment.Hashbang, re.group('c1')
        yield re.start('opcode'), String.Doc, re.group('opcode')
        yield re.start('c2'), Comment.Hashbang, re.group('c2')
        yield re.start('rcode'), REPLY_CODES.get(re.group('rcode'), Name.Tag), re.group('rcode')
        yield re.start('c3'), Comment.Hashbang, re.group('c3')
        yield re.start('nonce'), String.Doc, re.group('nonce')
        yield re.start('c4'), Comment.Hashbang, re.group('c4')
        yield re.start('flags'), String.Doc, re.group('flags')
        yield re.start('c5'), Comment.Hashbang, re.group('c5')
        yield re.start('nquery'), String.Doc, re.group('nquery')
        yield re.start('c6'), Comment.Hashbang, re.group('c6')
        yield re.start('nanswer'), String.Doc, re.group('nanswer')
        yield re.start('c7'), Comment.Hashbang, re.group('c7')
        yield re.start('nauth'), String.Doc, re.group('nauth')
        yield re.start('c8'), Comment.Hashbang, re.group('c8')
        yield re.start('nadd'), String.Doc, re.group('nadd')
        ctx.pos = re.end()

    tokens = {
        'root': [
            (r'(?P<c1>; <<>>)(?P<version>.*)(?P<c2><<>>)(?P<query>.*\n)', parse_dig_header),
            (r'(?P<c1>;; ->>HEADER<<- opcode: )(?P<opcode>[^\s]+)(?P<c2>, status: )(?P<rcode>[^\s]+)(?P<c3>, id: )(?P<nonce>[^\s]+\n)(?P<c4>;; flags: )(?P<flags>[^;]+)(?P<c5>; QUERY: )(?P<nquery>\d+)(?P<c6>, ANSWER: )(?P<nanswer>\d+)(?P<c7>, AUTHORITY: )(?P<nauth>\d+)(?P<c8>, ADDITIONAL: )(?P<nadd>\d+\n)', parse_response_header),
            (r';;.*', Generic.Subheading),  # comment section headers
            (r';.*', Comment.Single),  # comment section headers
            (r'\n', Generic),
            (r'^', Generic, 'rname')
        ],
        'rname': [
            (r'\s', Generic, 'wsp1'),
            (r'\\[0-7]{3}', String.Escape),
            (r'\\.', String.Escape),
            (r'.', Name.Function)
        ],
        'wsp1': [
            (r'\s', Generic),
            (r'[0-9]+', Number.Integer, 'wsp2'),
        ],
        'wsp2': [
            (r'\s', Generic),
            (r'[^\s]+', Name.Class, 'wsp3'),
        ],
        'wsp3': [
            (r'\s', Generic),
            (r'[^\s]+', Name.Tag, 'value'),
        ],
        'value': [
            (r';.*$', Comment.Special),
            (r'\(', Punctuation, 'multiline'),
            (r'\n', Generic, '#pop:5'),
            (r'\s', Generic),
            (r'"', String.Double, 'string'),
            (r'.', String.Single)
        ],
        'string': [
            (r'\\[0-7]{3}', String.Escape),
            (r'\\.', String.Escape),
            (r'"', String.Double, '#pop'),
            (r'.', String.Double)
        ],
        'multiline': [
            (r';.*$', Comment.Special),
            (r'\)', Punctuation, '#pop'),
            (r'\s', Generic),
            (r'"', String.Double, 'string'),
            (r'.', String.Single)
        ]
    }