import os
import struct
import sys
import json
import time
from tqdm import tqdm

IMAGE_EXTENSIONS = ['.stex', '.ctex']
AUDIO_EXTENSIONS = ['.oggstr', '.mp3str']

EXPORT_MAP = {
    '.oggstr': '.ogg',
    '.mp3str': '.mp3',
}

PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
IEND_CHUNK = b"IEND"
RIFF_SIGNATURE = b"RIFF"
WEBP_FOURCC = b"WEBP"

def log_info(msg):
    print(f"\033[94m[INFO]\033[0m {msg}")

def log_warn(msg):
    print(f"\033[93m[WARN]\033[0m {msg}")

def log_error(msg):
    print(f"\033[91m[ERROR]\033[0m {msg}")

def read_binary_resource(path):
    with open(path, "rb") as f:
        return f.read()

def write_binary_resource(path, data):
    with open(path, "wb") as f:
        f.write(data)

def find_all_webp_images(data):
    positions = []
    i = 0
    data_len = len(data)
    MAX_SIZE = 50 * 1024 * 1024

    while i <= data_len - 12:
        if data[i:i+4] == RIFF_SIGNATURE and data[i+8:i+12] == WEBP_FOURCC:
            size = struct.unpack("<I", data[i+4:i+8])[0]
            total_len = 8 + size
            if total_len < 20 or total_len > MAX_SIZE:
                i += 1
                continue
            if i + total_len > data_len:
                i += 1
                continue
            positions.append((i, total_len))
            i += total_len
        else:
            i += 1
    return positions

def extract_all_images_to_folder(filepath, output_folder):
    data = read_binary_resource(filepath)
    if not os.path.isdir(output_folder):
        os.makedirs(output_folder)

    if PNG_SIGNATURE in data:
        start = data.find(PNG_SIGNATURE)
        iend = data.find(IEND_CHUNK, start)
        if iend == -1:
            log_warn(f"IEND chunk not found in {filepath}")
            return False
        end = iend + 8

        png_data = data[start:end]
        out_path = os.path.join(output_folder, "image_1.png")
        with open(out_path, "wb") as f:
            f.write(png_data)

        offsets = [{"start": start, "length": len(png_data), "filename": "image_1.png"}]
        with open(os.path.join(output_folder, "offsets.json"), "w") as f:
            json.dump(offsets, f, indent=2)

        log_info(f"Extracted PNG from {filepath} to {out_path}")
        return True

    elif RIFF_SIGNATURE in data and WEBP_FOURCC in data:
        positions = find_all_webp_images(data)
        if not positions:
            log_warn(f"No WebP images found in {filepath}")
            return False

        offsets = []
        for i, (start, length) in enumerate(positions, 1):
            out_path = os.path.join(output_folder, f"image_{i}.webp")
            with open(out_path, "wb") as f:
                f.write(data[start:start+length])
            offsets.append({"start": start, "length": length, "filename": f"image_{i}.webp"})
            log_info(f"Extracted WebP #{i} from {filepath} to {out_path}")

        with open(os.path.join(output_folder, "offsets.json"), "w") as f:
            json.dump(offsets, f, indent=2)

        return True
    else:
        log_warn(f"No known image signature found in {filepath}")
        return False

def patch_all_images_from_folder(resource_path, input_folder, output_path):
    data = bytearray(read_binary_resource(resource_path))
    offsets_file = os.path.join(input_folder, "offsets.json")
    if not os.path.isfile(offsets_file):
        log_error(f"Missing offsets.json in {input_folder}")
        return False

    with open(offsets_file, "r") as f:
        offsets = json.load(f)

    sidecar_images = sorted(f for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.webp')))
    if len(sidecar_images) != len(offsets):
        log_error(f"Mismatch between sidecar images and offsets count")
        return False

    for entry in offsets:
        start = entry['start']
        length = entry['length']
        filename = entry['filename']
        img_path = os.path.join(input_folder, filename)
        if not os.path.isfile(img_path):
            log_error(f"Missing image file {filename} in {input_folder}")
            return False

        img_data = read_binary_resource(img_path)
        if len(img_data) > length:
            log_warn(f"Skipped patching {filename}: size {len(img_data)} larger than embedded slot {length}")
            continue  # Skip patching this image
        elif len(img_data) < length:
            img_data += b'\x00' * (length - len(img_data))

        data[start:start + length] = img_data
        log_info(f"Patched {filename} at offset {start}")

    write_binary_resource(output_path, data)
    log_info(f"Saved patched file: {output_path}")
    return True

def extract_embedded_data(resource_path, output_path_or_folder):
    if os.path.splitext(output_path_or_folder)[1]:
        output_folder = os.path.splitext(output_path_or_folder)[0] + "_export"
    else:
        output_folder = output_path_or_folder
    return extract_all_images_to_folder(resource_path, output_folder)

def patch_embedded_data(resource_path, new_file_path_or_folder, output_resource_path):
    input_folder = new_file_path_or_folder
    if os.path.isfile(new_file_path_or_folder):
        input_folder = os.path.splitext(new_file_path_or_folder)[0] + "_export"
    ext = os.path.splitext(resource_path)[1].lower()
    if ext in AUDIO_EXTENSIONS:
        return patch_audio_data(resource_path, new_file_path_or_folder, output_resource_path)
    return patch_all_images_from_folder(resource_path, input_folder, output_resource_path)

def find_audio_section(data):
    signatures = [b"OggS", b"ID3", b"RIFF"]
    positions = [data.find(sig) for sig in signatures if data.find(sig) != -1]
    if not positions:
        return None
    start = min(positions)
    length = len(data) - start
    return start, length

def extract_audio_data(resource_path, output_path):
    data = read_binary_resource(resource_path)
    chunk = find_audio_section(data)
    if chunk is None:
        log_warn(f"No embedded audio signature found in {resource_path}")
        return False
    start, length = chunk
    embedded = data[start:start + length]
    with open(output_path, "wb") as f:
        f.write(embedded)
    log_info(f"Extracted audio from {resource_path} to {output_path}")
    return True

def patch_audio_data(resource_path, new_audio_path, output_resource_path):
    data = bytearray(read_binary_resource(resource_path))
    chunk = find_audio_section(data)
    if chunk is None:
        log_warn(f"No embedded audio signature found in {resource_path}")
        return False
    start, length = chunk
    new_data = read_binary_resource(new_audio_path)
    if len(new_data) > length:
        log_warn(f"Skipped patching audio {os.path.basename(new_audio_path)}: size {len(new_data)} > embedded slot {length}")
        return False
    data[start:start+len(new_data)] = new_data
    if len(new_data) < length:
        data[start+len(new_data):start+length] = b'\x00' * (length - len(new_data))
        log_info(f"Padded new audio data with zeros")
    write_binary_resource(output_resource_path, data)
    log_info(f"Patched audio in {resource_path}")
    return True

def export_folder(folder):
    candidate_files = []
    for root, _, files in os.walk(folder):
        for f in files:
            ext = os.path.splitext(f)[1].lower()
            if ext in IMAGE_EXTENSIONS + AUDIO_EXTENSIONS:
                candidate_files.append(os.path.join(root, f))
    if not candidate_files:
        log_warn("No resources found for export.")
        return

    log_info(f"Exporting {len(candidate_files)} resources...")
    for path in tqdm(candidate_files, unit="file"):
        root = os.path.dirname(path)
        base = os.path.splitext(os.path.basename(path))[0]
        ext = os.path.splitext(path)[1].lower()

        try:
            if ext in AUDIO_EXTENSIONS:
                out_file = os.path.join(root, base + EXPORT_MAP.get(ext, ""))
                if os.path.exists(out_file):
                    log_warn(f"Skipping existing {out_file}")
                    continue
                extract_audio_data(path, out_file)

            else:
                out_dir = os.path.join(root, base + "_export")
                if not os.path.exists(out_dir):
                    os.makedirs(out_dir)
                extract_all_images_to_folder(path, out_dir)
        except Exception as e:
            log_error(f"Failed to export {path}: {e}")

def import_folder(folder):
    candidate_files = []
    for root, _, files in os.walk(folder):
        for f in files:
            ext = os.path.splitext(f)[1].lower()
            if ext in IMAGE_EXTENSIONS + AUDIO_EXTENSIONS:
                candidate_files.append(os.path.join(root, f))
    if not candidate_files:
        log_warn("No resources found for import.")
        return

    log_info(f"Importing {len(candidate_files)} resources...")
    for path in tqdm(candidate_files, unit="file"):
        root = os.path.dirname(path)
        base = os.path.splitext(os.path.basename(path))[0]
        ext = os.path.splitext(path)[1].lower()

        try:
            if ext in AUDIO_EXTENSIONS:
                sidecar = os.path.join(root, base + EXPORT_MAP.get(ext, ""))
                if not os.path.isfile(sidecar):
                    log_warn(f"Audio sidecar missing: {sidecar}")
                    continue
                if patch_audio_data(path, sidecar, path):
                    os.remove(sidecar)
                    log_info(f"Patched and removed audio sidecar {sidecar}")
                else:
                    log_warn(f"Failed to patch audio {path}")

            else:
                sidecar_dir = os.path.join(root, base + "_export")
                if not os.path.isdir(sidecar_dir):
                    log_warn(f"Image sidecar folder missing: {sidecar_dir}")
                    continue
                if patch_all_images_from_folder(path, sidecar_dir, path):
                    for f in os.listdir(sidecar_dir):
                        try:
                            os.remove(os.path.join(sidecar_dir, f))
                        except:
                            pass
                    try:
                        os.rmdir(sidecar_dir)
                    except:
                        pass
                    log_info(f"Patched and removed image sidecar folder {sidecar_dir}")
                else:
                    log_warn(f"Failed to patch image {path}")
        except Exception as e:
            log_error(f"Error importing {path}: {e}")

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=["export", "import"])
    parser.add_argument("folder")
    args = parser.parse_args()

    if not os.path.isdir(args.folder):
        log_error(f"Folder not found: {args.folder}")
        sys.exit(1)

    start_time = time.time()

    if args.mode == "export":
        export_folder(args.folder)
    else:
        import_folder(args.folder)

    end_time = time.time()
    elapsed = end_time - start_time
    log_info(f"Elapsed time: {elapsed:.2f} seconds")


if __name__ == "__main__":
    main()
-------------------------------------------
For first file where no known image signature fould message appear empty folder BackButton.svg-b32a237cd77a016d1ee8e2b5f3d77850.bptc_export isn't being deleted
G:\GDRE_tools-v1.0.0-rc.3-windows>G:\GDRE_tools-v1.0.0-rc.3-windows\dist\godot\godot.exe export "I:\games\Honest_Bond\UAGC\.godot\imported"
[INFO] Exporting 2 resources...
  0%|                                                                                          | 0/2 [00:00<?, ?file/s][WARN] No known image signature found in I:\games\Honest_Bond\UAGC\.godot\imported\BackButton.svg-b32a237cd77a016d1ee8e2b5f3d77850.bptc.ctex
[INFO] Extracted PNG from I:\games\Honest_Bond\UAGC\.godot\imported\BackButton_Simple.svg-0fee8d161e40c3123d63a5100e1fc8ac.ctex to I:\games\Honest_Bond\UAGC\.godot\imported\BackButton_Simple.svg-0fee8d161e40c3123d63a5100e1fc8ac_export\image_1.png
100%|█████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 198.80file/s]
[INFO] Elapsed time: 0.02 seconds