"""
Serialization functions for the SimpleMRS format.
"""

from pathlib import Path
from typing import Optional

from delphin.util import Lexer
from delphin import predicate
from delphin.lnk import Lnk
from delphin.sembase import (role_priority, property_priority)
from delphin import variable
from delphin.mrs import (
    EP,
    HCons,
    ICons,
    MRS,
    MRSSyntaxError,
    CONSTANT_ROLE)


CODEC_INFO = {
    'representation': 'mrs',
}


TOP_FEATURE = 'TOP'


##############################################################################
##############################################################################
# Pickle-API methods


def load(source):
    """
    Deserialize SimpleMRSs from a file (handle or filename)

    Args:
        source (str, file): input filename or file object
    Returns:
        a list of MRS objects
    """
    if hasattr(source, 'read'):
        ms = list(_decode(source))
    else:
        source = Path(source).expanduser()
        with source.open() as fh:
            ms = list(_decode(fh))
    return ms


def loads(s):
    """
    Deserialize SimpleMRS string representations

    Args:
        s (str): a SimpleMRS string
    Returns:
        a list of MRS objects
    """
    ms = list(_decode(s.splitlines()))
    return ms


def dump(ms, destination, properties=True, lnk=True,
         indent=False, encoding='utf-8'):
    """
    Serialize MRS objects to SimpleMRS and write to a file

    Args:
        ms: an iterator of MRS objects to serialize
        destination: filename or file object where data will be written
        properties: if `False`, suppress morphosemantic properties
        lnk: if `False`, suppress surface alignments and strings
        indent (bool, int): if `True` or an integer value, add
            newlines and indentation
        encoding (str): if *destination* is a filename, write to the
            file with the given encoding; otherwise it is ignored
    """
    text = dumps(ms, properties=properties, lnk=lnk, indent=indent)
    if hasattr(destination, 'write'):
        print(text, file=destination)
    else:
        destination = Path(destination).expanduser()
        with destination.open('w', encoding=encoding) as fh:
            print(text, file=fh)


def dumps(ms, properties=True, lnk=True, indent=False):
    """
    Serialize MRS objects to a SimpleMRS representation

    Args:
        ms: an iterator of MRS objects to serialize
        properties: if `False`, suppress variable properties
        lnk: if `False`, suppress surface alignments and strings
        indent (bool, int): if `True` or an integer value, add
            newlines and indentation
    Returns:
        a SimpleMRS string representation of a corpus of MRS objects
    """
    return _encode(ms, properties, lnk, indent)


def decode(s):
    """
    Deserialize an MRS object from a SimpleMRS string.
    """
    lexer = SimpleMRSLexer.lex(s.splitlines())
    return _decode_mrs(lexer)


def encode(m, properties=True, lnk=True, indent=False):
    """
    Serialize a MRS object to a SimpleMRS string.

    Args:
        m: an MRS object
        properties (bool): if `False`, suppress variable properties
        lnk: if `False`, suppress surface alignments and strings
        indent (bool, int): if `True` or an integer value, add
            newlines and indentation
    Returns:
        a SimpleMRS-serialization of the MRS object
    """
    return _encode([m], properties, lnk, indent)


##############################################################################
##############################################################################
# Deserialization


SimpleMRSLexer = Lexer(
    tokens=[
        (r'\[', 'LBRACK:['),
        (r'\]', 'RBRACK:]'),
        (r'<(?:-?\d+[:#]-?\d+'
         r'|@\d+'
         r'|\d+(?: +\d+)*)>', 'LNK:a lnk value'),
        (r'"([^"\\]*(?:\\.[^"\\]*)*)"', 'DQSTRING:a string'),
        (r"'([^ \n:<>\[\]])", 'SQSYMBOL:a quoted symbol'),
        (r'_[^\s_]+'  # lemma
         r'_[nvajrscpqxud]'  # pos
         r'(?:_(?:[^\s_<]|<(?![-0-9:#@ ]*>\s))+)?'  # optional sense
         r'(?:_rel)?',  # optional suffix
         'PREDICATE:a surface predicate'),
        (r'<', 'LANGLE:<'),
        (r'>', 'RANGLE:>'),
        (r'([^\s:<>\[\]]+):', 'FEATURE:a feature'),
        (r'(?:[^ \n\]<]+'
         r'|<(?![-0-9:#@ ]*>\s))+', 'SYMBOL:a symbol'),
        (r'[^\s]', 'UNEXPECTED'),
    ],
    error_class=MRSSyntaxError)


LBRACK    = SimpleMRSLexer.tokentypes.LBRACK
RBRACK    = SimpleMRSLexer.tokentypes.RBRACK
LNK       = SimpleMRSLexer.tokentypes.LNK
DQSTRING  = SimpleMRSLexer.tokentypes.DQSTRING
SQSYMBOL  = SimpleMRSLexer.tokentypes.SQSYMBOL
PREDICATE = SimpleMRSLexer.tokentypes.PREDICATE
LANGLE    = SimpleMRSLexer.tokentypes.LANGLE
RANGLE    = SimpleMRSLexer.tokentypes.RANGLE
FEATURE   = SimpleMRSLexer.tokentypes.FEATURE
SYMBOL    = SimpleMRSLexer.tokentypes.SYMBOL


def _decode(lineiter):
    lexer = SimpleMRSLexer.lex(lineiter)
    try:
        while lexer.peek():
            yield _decode_mrs(lexer)
    except StopIteration:
        pass


def _decode_mrs(lexer):
    top = index = lnk = surface = identifier = None
    rels = []
    hcons = []
    icons = []
    variables = {}
    lexer.expect_type(LBRACK)
    lnk = _decode_lnk(lexer)
    surface = _decode_dqstring(lexer.accept_type(DQSTRING))
    feature = lexer.accept_type(FEATURE)
    while feature is not None:
        feature = feature.upper()
        if feature in ('LTOP', 'TOP'):
            top = lexer.expect_type(SYMBOL).lower()
        elif feature == 'INDEX':
            index = _decode_variable(lexer, variables)
        elif feature == 'RELS':
            lexer.expect_type(LANGLE)
            while lexer.peek()[0] == LBRACK:
                rels.append(_decode_rel(lexer, variables))
            lexer.expect_type(RANGLE)
        elif feature == 'HCONS':
            lexer.expect_type(LANGLE)
            while lexer.peek()[0] == SYMBOL:
                hcons.append(_decode_cons(lexer, HCons, variables))
            lexer.expect_type(RANGLE)
        elif feature == 'ICONS':
            lexer.expect_type(LANGLE)
            while lexer.peek()[0] == SYMBOL:
                icons.append(_decode_cons(lexer, ICons, variables))
            lexer.expect_type(RANGLE)
        else:
            raise ValueError('invalid feature: ' + feature)
        feature = lexer.accept_type(FEATURE)
    lexer.expect_type(RBRACK)
    return MRS(top, index, rels, hcons,
               icons=icons, variables=variables,
               lnk=lnk, surface=surface, identifier=identifier)


def _decode_lnk(lexer):
    lnk = lexer.accept_type(LNK)
    if lnk is not None:
        lnk = Lnk(lnk)
    return lnk


def _decode_dqstring(dqstring: Optional[str]) -> Optional[str]:
    if dqstring is not None:
        dqstring = _unescape(dqstring)
    return dqstring


def _decode_variable(lexer, variables):
    var = lexer.expect_type(SYMBOL).lower()
    if var not in variables:
        variables[var] = {}
    props = variables[var]
    if lexer.accept_type(LBRACK):
        lexer.accept_type(SYMBOL)  # variable type
        feature = lexer.accept_type(FEATURE)
        while feature is not None:
            value = lexer.expect_type(SYMBOL)
            props[feature.upper()] = value.lower()
            feature = lexer.accept_type(FEATURE)
        lexer.expect_type(RBRACK)
    return var


def _decode_rel(lexer, variables):
    args = {}
    surface = None
    lexer.expect_type(LBRACK)
    pred = _decode_predicate(lexer)
    lnk = _decode_lnk(lexer)
    surface = _decode_dqstring(lexer.accept_type(DQSTRING))
    _, label = lexer.expect((FEATURE, 'LBL'), (SYMBOL, None))
    # any remaining are arguments or a constant
    role = lexer.accept_type(FEATURE)
    while role is not None:
        role = role.upper()
        if role == 'CARG':
            value = _decode_dqstring(lexer.expect_type(DQSTRING))
        else:
            value = _decode_variable(lexer, variables)
        args[role] = value
        role = lexer.accept_type(FEATURE)
    lexer.expect_type(RBRACK)
    return EP(pred,
              label.lower(),
              args=args,
              lnk=lnk,
              surface=surface,
              base=None)


def _decode_predicate(lexer) -> str:
    predstring = lexer.accept_type(DQSTRING)
    if predstring is not None:
        predstring = _decode_dqstring(predstring)
    else:
        predstring = lexer.choice_type(SQSYMBOL, PREDICATE, SYMBOL)[1]
    return predicate.normalize(predstring)


def _decode_cons(lexer, cls, variables):
    lhs = _decode_variable(lexer, variables)
    relation = lexer.expect_type(SYMBOL).lower()
    rhs = _decode_variable(lexer, variables)
    return cls(lhs, relation, rhs)


##############################################################################
##############################################################################
# Encoding

def _encode(ms, properties, lnk, indent):
    if indent is None or indent is False:
        indent = False  # normalize None to False
        delim = ' '
    else:
        indent = True  # normalize integers to True
        delim = '\n'
    return delim.join(_encode_mrs(m, properties, lnk, indent) for m in ms)


def _encode_mrs(m, properties, lnk, indent):
    delim = '\n  ' if indent else ' '
    if properties:
        varprops = dict(m.variables)
    else:
        varprops = {}
    parts = [
        _encode_surface_info(m, lnk),
        _encode_hook(m, varprops, indent),
        _encode_rels(m.rels, varprops, lnk, indent),
        _encode_hcons(m.hcons),
        _encode_icons(m.icons, varprops)
    ]
    return '[ {} ]'.format(
        delim.join(
            ' '.join(tokens) for tokens in parts if tokens))


def _encode_surface_info(m, lnk):
    tokens = []
    if lnk:
        if m.lnk:
            tokens.append(str(m.lnk))
        if m.surface is not None:
            tokens.append('"{}"'.format(_escape(m.surface)))
    return tokens


def _encode_hook(m, varprops, indent):
    delim = '\n  ' if indent else ' '
    tokens = []
    if m.top is not None:
        tokens.append('{}: {}'.format(TOP_FEATURE, m.top))
    if m.index is not None:
        tokens.append('INDEX: {}'.format(_encode_variable(m.index, varprops)))
    if tokens:
        tokens = [delim.join(tokens)]
    return tokens


def _encode_variable(var, varprops):
    tokens = [var]
    if varprops.get(var):
        tokens.append('[')
        tokens.append(variable.type(var))
        for prop in sorted(varprops[var], key=property_priority):
            val = varprops[var][prop]
            tokens.append(prop + ':')
            tokens.append(val)
        tokens.append(']')
        del varprops[var]
    return ' '.join(tokens)


def _encode_rels(rels, varprops, lnk, indent):
    delim = ('\n  ' + ' ' * len('RELS: < ')) if indent else ' '
    tokens = []
    for rel in rels:
        pred = rel.predicate
        if lnk:
            pred += str(rel.lnk)
        reltoks = ['[', pred]
        if lnk and rel.surface is not None:
            reltoks.append('"{}"'.format(_escape(rel.surface)))
        reltoks.extend(('LBL:', rel.label))
        for role in sorted(rel.args, key=role_priority):
            arg = rel.args[role]
            if role == CONSTANT_ROLE:
                arg = '"{}"'.format(_escape(arg))
            else:
                arg = _encode_variable(arg, varprops)
            reltoks.extend((role + ':', arg))
        reltoks.append(']')
        tokens.append(' '.join(reltoks))
    if tokens:
        tokens = ['RELS: <'] + [delim.join(tokens)] + ['>']
    return tokens


def _encode_hcons(hcons):
    tokens = ['{} {} {}'.format(hc.hi, hc.relation, hc.lo)
              for hc in hcons]
    if tokens:
        tokens = ['HCONS: <'] + [' '.join(tokens)] + ['>']
    return tokens


def _encode_icons(icons, varprops):
    tokens = ['{} {} {}'.format(_encode_variable(ic.left, varprops),
                                ic.relation,
                                _encode_variable(ic.right, varprops))
              for ic in icons]
    if tokens:
        tokens = ['ICONS: <'] + [' '.join(tokens)] + ['>']
    return tokens


# Character Escaping


_ESCAPES = {
    '\\': '\\\\',
    '"': '\\"',
}


_UNESCAPES = {
    '\\\\': '\\',
    '\\"': '"',
}


def _escape(s: str) -> str:
    return "".join(_ESCAPES.get(c, c) for c in s)


def _unescape(s: str) -> str:
    if not s:
        return s
    cs = []
    i = 0
    while i < len(s):
        if s[i] == '\\' and (i + 1) < len(s):
            cs.append(s[i+1])
            i += 2
        else:
            cs.append(s[i])
            i += 1
    return "".join(cs)
