"""REDengine network protocol, adapted from MIT-licensed rusty_witcher_debugger.

Reference: SpontanCombust/rusty_witcher_debugger at
5a0cf7da77af512d89ac4db5dd7fb38a1aeef7f3. Remastered uses UTF-8 in places
where older builds used UTF-16. No keyboard, clipboard or settings transport.
"""
import socket
import struct
import time


def encode_packet(values):
    payload = bytearray()
    for value in values:
        if isinstance(value, str):
            data = value.encode('utf-8')
            if len(data) > 32767:
                raise ValueError('Protocol string too long')
            payload.extend(b'\xac\x08\x81\x16' + struct.pack('>H', len(data)) + data)
        elif isinstance(value, int):
            payload.extend(b'\x81\x32' + struct.pack('>I', value & 0xffffffff))
        else:
            raise TypeError('Unsupported protocol value')
    size = len(payload) + 6
    if size > 65535:
        raise ValueError('Protocol packet too large')
    return b'\xde\xad' + struct.pack('>H', size) + payload + b'\xbe\xef'


def decode_payload(payload):
    values, pos = [], 0
    formats = {b'\x81\x08': '>b', b'\x81\x16': '>h', b'\x81\x32': '>i',
               b'\x71\x32': '>I', b'\x81\x64': '>q'}
    while pos < len(payload):
        tag = payload[pos:pos+2]; pos += 2
        if tag in formats:
            fmt = formats[tag]; n = struct.calcsize(fmt)
            if pos+n > len(payload): raise ValueError('Truncated number')
            values.append(struct.unpack_from(fmt, payload, pos)[0]); pos += n
        elif tag in (b'\xac\x08', b'\x9c\x16'):
            if payload[pos:pos+2] != b'\x81\x16' or pos+4 > len(payload):
                raise ValueError('Invalid string length tag')
            n = struct.unpack_from('>H', payload, pos+2)[0]; pos += 4
            if tag == b'\x9c\x16': n *= 2
            if pos+n > len(payload): raise ValueError('Truncated string')
            values.append(payload[pos:pos+n].decode('utf-8' if tag == b'\xac\x08' else 'utf-16-be'))
            pos += n
        else:
            raise ValueError(f'Unknown protocol tag {tag.hex()}')
    return values


class Connection:
    def __init__(self, timeout=5):
        self.socket = socket.create_connection(('127.0.0.1', 37001), timeout)
        self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.timeout = timeout
        self.deadline = time.monotonic() + timeout
        for namespace in ('Remote', 'ScriptCompiler', 'scripts', 'ScriptDebugger', 'Config'):
            self.send(['BIND', namespace])

    def send(self, values):
        self.socket.sendall(encode_packet(values))

    def read_exact(self, n):
        data = bytearray()
        while len(data) < n:
            left = self.deadline-time.monotonic()
            if left <= 0: raise TimeoutError('Network response deadline elapsed')
            self.socket.settimeout(left)
            chunk = self.socket.recv(n-len(data))
            if not chunk: raise ConnectionError('Game closed the connection')
            data.extend(chunk)
        return bytes(data)

    def receive(self):
        head = self.read_exact(4)
        if head[:2] != b'\xde\xad': raise ValueError('Invalid packet head')
        size = struct.unpack('>H', head[2:])[0]
        if not 6 <= size <= 65535: raise ValueError('Invalid packet size')
        body = self.read_exact(size-4)
        if body[-2:] != b'\xbe\xef': raise ValueError('Invalid packet tail')
        return decode_payload(body[:-2])

    def request(self, values, prefix, markers=None):
        self.deadline = time.monotonic()+self.timeout
        self.send(values)
        packets = []
        while True:
            packet = self.receive(); packets.append(packet)
            if packet[:len(prefix)] == prefix and (not markers or
                    len(packet)>len(prefix) and isinstance(packet[len(prefix)],str) and
                    any(marker in packet[len(prefix)] for marker in markers)):
                return packet, packets
            if len(packets) > 4096: raise ValueError('Too many unrelated packets')

    def execute(self, expression, markers=None):
        return self.request(['Remote', 0x12345678, 0x81160008, expression],
                            [0x12345678, -2129264632], markers)

    def recompile(self):
        self.deadline = time.monotonic()+self.timeout
        self.send(['scripts', 'reload'])
        packets = []
        while True:
            packet = self.receive(); packets.append(packet)
            if packet[:2] == ['ScriptCompiler', 'finished']:
                return {'ok': packet[2] == 0 and not any(p[:2] == ['ScriptCompiler', 'error'] for p in packets),
                        'packets': packets}
            if len(packets) > 20000: raise ValueError('Compiler output limit exceeded')

    def close(self):
        self.socket.close()

    def __enter__(self): return self
    def __exit__(self, *_): self.close()