"""Translate only direct display-call literals in Dragon .dsc scripts.

Source defaults to the original single-byte encoding; JSON mappings are UTF-8
and use decoded string values (including actual control characters). The game
output encodes translated literals as GB2312 with script escapes restored.
"""
import argparse
from dataclasses import dataclass
import hashlib
import json
from pathlib import Path
import re


ALLOWED = {'DisplayMessage': 2, 'SetMissionDescription': 1}
ESCAPES = {ord(k): ord(v) for k, v in {'a': '\a', 'b': '\b', 'f': '\f',
           'n': '\n', 'r': '\r', 't': '\t', 'v': '\v', '\\': '\\', '"': '"'}.items()}
ENCODE_ESCAPES = {v: bytes((92, k)) for k, v in ESCAPES.items()}
FORMAT = re.compile(r'%%|%[-+ #0]*[\d*]*(?:\.[\d*]+)?[hlL]?[diouxXeEfgGcs%]')
CONTROLS = re.compile(r'[\a\b\f\n\r\t\v]')


@dataclass(frozen=True)
class Literal:
    start: int
    end: int
    text: str
    call: str
    argument: int


@dataclass(frozen=True)
class Token:
    kind: str
    start: int
    end: int
    value: str


def decode_literal(raw: bytes, encoding: str) -> str:
    # A continuation includes the physical newline and indentation, not a
    # game newline. Original scripts also contain backslash-space-CRLF.
    raw = re.sub(rb'\\[ \t]*\r?\n[ \t]*', b'', raw)
    out = bytearray()
    i = 0
    while i < len(raw):
        if raw[i] == 92:
            if i + 1 == len(raw) or raw[i + 1] not in ESCAPES:
                raise ValueError('unsupported script escape in string literal')
            out.append(ESCAPES[raw[i + 1]])
            i += 2
        else:
            out.append(raw[i])
            i += 1
    return out.decode(encoding)


def tokenize(source: bytes, encoding: str) -> list[Token]:
    tokens = []
    i = 0
    while i < len(source):
        b = source[i]
        if b in b' \t\r\n':
            i += 1
        elif b == 35 or source[i:i+2] == b'//':
            end = source.find(b'\n', i)
            i = len(source) if end < 0 else end + 1
        elif source[i:i+2] == b'/*':
            end = source.find(b'*/', i + 2)
            if end < 0:
                raise ValueError('unterminated block comment')
            i = end + 2
        elif b == 34:
            start = i
            i += 1
            while i < len(source) and source[i] != 34:
                i += 2 if source[i] == 92 else 1
            if i >= len(source):
                raise ValueError('unterminated string at byte %d' % start)
            i += 1
            tokens.append(Token('string', start, i, decode_literal(source[start+1:i-1], encoding)))
        elif b in b'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_':
            start = i
            i += 1
            while i < len(source) and source[i] in b'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_0123456789':
                i += 1
            tokens.append(Token('name', start, i, source[start:i].decode('ascii')))
        else:
            tokens.append(Token('symbol', i, i + 1, chr(b)))
            i += 1
    return tokens


def extract_display_literals(source: bytes, encoding: str = 'latin1') -> list[Literal]:
    tokens = tokenize(source, encoding)
    found = []
    for index, token in enumerate(tokens[:-1]):
        if token.kind != 'name' or token.value not in ALLOWED or tokens[index+1].value != '(':
            continue
        arguments = [[]]
        nesting = ['(']
        for current in tokens[index+2:]:
            value = current.value if current.kind == 'symbol' else None
            if value in ('(', '[', '{'):
                nesting.append(value)
            elif value in (')', ']', '}'):
                if not nesting or nesting[-1] != {')': '(', ']': '[', '}': '{'}[value]:
                    raise ValueError('unbalanced display call at byte %d' % token.start)
                nesting.pop()
                if not nesting:
                    break
            if value == ',' and len(nesting) == 1:
                arguments.append([])
            else:
                arguments[-1].append(current)
        else:
            raise ValueError('unterminated display call at byte %d' % token.start)
        if len(arguments) != ALLOWED[token.value]:
            continue
        for number, arg in enumerate(arguments):
            if len(arg) == 1 and arg[0].kind == 'string':
                lit = arg[0]
                found.append(Literal(lit.start, lit.end, lit.value, token.value, number))
    return sorted(found, key=lambda literal: literal.start)


def encode_literal(text: str) -> bytes:
    raw = text.encode('gb2312')
    return b'"' + b''.join(ENCODE_ESCAPES.get(b, bytes((b,))) for b in raw) + b'"'


def patch_display_calls(source: bytes, translations: dict[str, str]) -> bytes:
    literals = extract_display_literals(source)
    missing = set(translations) - {literal.text for literal in literals}
    if missing:
        raise ValueError('translation source not found: %r' % sorted(missing))
    replacements = []
    for literal in literals:
        if literal.text not in translations:
            continue
        translated = translations[literal.text]
        if not isinstance(translated, str):
            raise ValueError('translation must be a string')
        if CONTROLS.findall(literal.text) != CONTROLS.findall(translated):
            raise ValueError('control sequence changed: %r' % literal.text)
        if FORMAT.findall(literal.text) != FORMAT.findall(translated):
            raise ValueError('format sequence changed: %r' % literal.text)
        replacements.append((literal.start, literal.end, encode_literal(translated)))
    result = source
    for start, end, value in reversed(replacements):
        result = result[:start] + value + result[end:]
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', required=True, type=Path)
    parser.add_argument('--mapping', required=True, type=Path)
    parser.add_argument('--out', required=True, type=Path)
    args = parser.parse_args()
    if args.source.resolve() == args.out.resolve():
        parser.error('output must differ from source')
    source = args.source.read_bytes()
    mapping = json.loads(args.mapping.read_text(encoding='utf-8'))
    if not isinstance(mapping, dict):
        parser.error('mapping must be an object of original text to translated text')
    result = patch_display_calls(source, mapping)
    changed = [literal for literal in extract_display_literals(source) if literal.text in mapping]
    manifest = {'source': str(args.source.resolve()), 'output': str(args.out.resolve()),
                'source_sha256': hashlib.sha256(source).hexdigest(),
                'output_sha256': hashlib.sha256(result).hexdigest(),
                'mapping_count': len(mapping), 'changed_literals': len(changed),
                'changes': [{'start': lit.start, 'end': lit.end, 'call': lit.call,
                             'argument': lit.argument} for lit in changed]}
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_bytes(result)
    args.out.with_suffix('.manifest.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')
    print('Translated %d literals using %d mappings' % (len(changed), len(mapping)))
    print('Source SHA256: ' + manifest['source_sha256'])
    print('Output SHA256: ' + manifest['output_sha256'])


if __name__ == '__main__':
    main()