import os
import json
import hashlib
import zlib
import struct
import sys
import re

# Decryption Keys (Hardcoded as per legacy logic)
REALKEY = bytearray([0x3C, 0xB5, 0x3C, 0x7F, 0x83, 0x94, 0xBA, 0x3B, 0x2B, 0xB2, 0x73, 0x5B, 0xEF, 0xEE, 0xE2, 0xA3, 0x3B, 0x2B, 0xCC, 0x66, 0x3D, 0xE5, 0x2C, 0xD7, 0x4D, 0x2E, 0x17, 0xE6, 0xF3])
REALKEYXXXX = bytearray([0x5F, 0xFA, 0xE7, 0xBA, 0xCC, 0xFE, 0xFB, 0x5C, 0x1A, 0xFB, 0xBD, 0xBB, 0x93, 0xB5, 0x83, 0xE7])
REALKEY2 = bytearray([0x1B, 0xC3, 0xAE, 0xF5, 0x87, 0x8D, 0xAF, 0x3F, 0x2B, 0xC2, 0xD3, 0xFC, 0xFE, 0xE6, 0xF3, 0xA1, 0x3C, 0x3C, 0xFC, 0xB4, 0x65])

# --- Encryption Utils ---

def xor(data, key):
    res = bytearray(len(data))
    v11 = 0
    for i in range(len(data)):
        res[i] = data[i] ^ key[v11]
        v11 = 7 if v11 == 28 else v11 + 1
    return res

def xor2(data, key):
    res = bytearray(len(data))
    v16 = 1
    res[0] = 68 # 'D'
    for i in range(1, 16):
        if v16 == 21: v16 = 7
        res[i] = data[i] ^ key[v16]
        v16 += 1
    return res

def xxtea_encrypt(data, key):
    # Standard XXTEA with length suffix
    if len(key) < 16:
        key = key + b'\0' * (16 - len(key))
    k = struct.unpack('<4I', key[:16])
    
    len_orig = len(data)
    num_uints = (len_orig >> 2)
    if (len_orig & 3) != 0: num_uints += 1
    
    # Pack data and add length suffix word
    v = [0] * (num_uints + 1)
    for i in range(len_orig):
        v[i >> 2] |= (data[i] << ((i & 3) << 3))
    v[num_uints] = len_orig # The length suffix
    
    n = len(v)
    if n < 2: return data
    
    DELTA, num6 = 0x9e3779B9, 0
    q = 6 + 52 // n
    z, y = v[n - 1], v[0]
    for _ in range(q):
        num6 = (num6 + DELTA) & 0xFFFFFFFF
        e = (num6 >> 2) & 3
        for p in range(n - 1):
            y = v[p + 1]
            mx = (((z >> 5 ^ y << 2) + (y >> 3 ^ z << 4)) ^ ((num6 ^ y) + (k[p & 3 ^ e] ^ z)))
            v[p] = (v[p] + mx) & 0xFFFFFFFF
            z = v[p]
        y, mx = v[0], (((z >> 5 ^ v[0] << 2) + (v[0] >> 3 ^ z << 4)) ^ ((num6 ^ v[0]) + (k[(n - 1) & 3 ^ e] ^ z)))
        v[n - 1] = (v[n - 1] + mx) & 0xFFFFFFFF
        z = v[n - 1]
    return struct.pack(f'<{n}I', *v)

def encrypt_file(input_path, output_path, outer_header=b'DHGAMES', inner_header=b'DHGAMES'):
    with open(input_path, "rb") as f: uncompressed = f.read()
    compressed = zlib.compress(uncompressed, level=6)
    block = inner_header + compressed
    key_bytes = xor2(REALKEYXXXX, REALKEY2)
    key = key_bytes.decode('utf-8', 'replace').encode('utf-8')
    
    encrypted_xxtea = xxtea_encrypt(block, key)
    xor_res = xor(encrypted_xxtea, REALKEY)
    final = outer_header + xor_res
    
    with open(output_path, "wb") as f: f.write(final)
    return final

def update_manifest(version_folder, rel_path, new_md5, new_size):
    manifest_path = os.path.join(version_folder, "manifest.json")
    if not os.path.exists(manifest_path):
        print(f"[!] Manifest not found at {manifest_path}")
        return False
    with open(manifest_path, "r") as f: manifest = json.load(f)
    found = False
    for entry in manifest:
        if entry["path"] == rel_path:
            entry["md5"], entry["size"] = new_md5, new_size
            found = True
            break
    if found:
        with open(manifest_path, "w") as f: json.dump(manifest, f, indent=2)
        print(f"[+] Updated manifest entry for {rel_path}")
        return True
    return False

def sync_project(new_manifest_hash):
    game_py = "handlers/game.py"
    if os.path.exists(game_py):
        with open(game_py, "r") as f: content = f.read()
        version_match = re.search(r'GAME_VERSION = "1\.34\.(\d+)"', content)
        if version_match:
            new_v = int(version_match.group(1)) + 1
            content = re.sub(r'GAME_VERSION = "1\.34\.\d+"', f'GAME_VERSION = "1.34.{new_v}"', content)
            print(f"[*] Updated GAME_VERSION to 1.34.{new_v}")
        content = re.sub(r'GAME_VERSION_HASH = "[0-9a-f]+"', f'GAME_VERSION_HASH = "{new_manifest_hash}"', content)
        with open(game_py, "w") as f: f.write(content)
        print(f"[*] Updated GAME_VERSION_HASH in {game_py}")

    server_py = "patch_server.py"
    if os.path.exists(server_py):
        with open(server_py, "r") as f: content = f.read()
        content = re.sub(r'hash_to_path\["[0-9a-f]+"\] = "manifest.json"', f'hash_to_path["{new_manifest_hash}"] = "manifest.json"', content)
        with open(server_py, "w") as f: f.write(content)
        print(f"[*] Updated manifest hash in {server_py}")

def get_existing_headers(path):
    """Peek into existing file to find exactly what headers it used."""
    if not os.path.exists(path):
        return b'DHGAMES', b'DHZAMES'
    try:
        with open(path, "rb") as f:
            data = f.read()
        outer = data[:7]
        # Decrypt to find inner
        xor_res = xor(data[7:], REALKEY)
        kb = xor2(REALKEYXXXX, REALKEY2).decode('utf-8', 'replace').encode('utf-8')
        dec = xxtea_decrypt(xor_res, kb)
        inner = dec[:7]
        return outer, inner
    except:
        return b'DHGAMES', b'DHZAMES'

def main():
    if len(sys.argv) < 5 or sys.argv[3].lower() != "with":
        print("Usage: python patch_tool.py <version_folder> <target_rel_path> with <source_file>")
        return

    version_folder = sys.argv[1]
    target_rel = sys.argv[2]
    source_file = sys.argv[4]

    print(f"[*] Patching: {target_rel}")
    
    target_full = os.path.join(version_folder, target_rel)
    outer_header, inner_header = get_existing_headers(target_full)
    
    print(f"[*] Logic: Dual-Header ({outer_header.decode()}/{inner_header.decode()}), Std XXTEA (+Length)")
    
    # 1. Encrypt
    encrypted_data = encrypt_file(source_file, target_full, outer_header, inner_header)
    
    # 2. Update Manifest
    new_md5 = hashlib.md5(encrypted_data).hexdigest()
    new_size = len(encrypted_data)
    update_manifest(version_folder, target_rel, new_md5, new_size)

    # 3. Calculate new manifest hash
    with open(os.path.join(version_folder, "manifest.json"), "rb") as f:
        new_manifest_hash = hashlib.md5(f.read()).hexdigest()

    print(f"\n[+] Patch Successful!")
    print(f"[*] New File MD5: {new_md5}")
    print(f"[*] New Manifest MD5: {new_manifest_hash}")
    
    sync_project(new_manifest_hash)
    print(f"[*] ACTION REQUIRED: Restart patch_server.py to apply changes.")

if __name__ == "__main__":
    main()