# Copyright (C) 2024 Thomas Hoppe (h0bB1T). All rights reserved.
#
# Unauthorized copying of this file via any medium is strictly prohibited.
# Proprietary and confidential.


import bpy, json, sys, os, hashlib, array, io, uuid

from typing import Any

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__))))


def load_cache(file: str) -> list[dict[str, Any]]:
    """
    Loads the asset cache file.
    """
    # If no cache exists yet, return nothing.
    if not os.path.exists(file):
        return []

    # We try to load the file and get the root list (libs).
    # If this fails, return empty list.
    try:
        with io.open(file, 'r', encoding='utf-8') as f:
            return json.load(f)['libs']
    except Exception as e:
        from shared import log_exception_shared
        log_exception_shared(e, context_msg=f'Failed to load asset cache: {file}')
        return []
    

def save_cache(file: str, data: dict[str, Any]):
    """
    Store modified cache file.
    """
    os.makedirs(os.path.dirname(file), exist_ok=True)

    tmp = file + '.' + str(uuid.uuid4()) + '.tmp'
    try:
        with io.open(tmp, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=4, ensure_ascii=False)
            f.flush()
            os.fsync(f.fileno())

        os.replace(tmp, file)
    except Exception as e:
        from shared import log_exception_shared
        log_exception_shared(e, context_msg=f'Failed to save asset cache: {file}')
        if os.path.exists(tmp):
            os.remove(tmp)
        return


def store_preview(cache_path: str, hash: str, data: list[float], size: tuple[int, int]):
    """
    Store preview image to cache folder.
    """
    # Dir must exist.
    os.makedirs(cache_path, exist_ok=True)

    # Build full name.
    full_name = os.path.join(cache_path, f'{hash}.png')

    # Check if already exists.
    if not os.path.exists(full_name):
        # No, create ..
        print(f'      Store preview: {full_name}')
        image = bpy.data.images.new('AW_TMP', alpha=False, width=size[0], height=size[1])
        try:
            image.filepath = full_name
            image.pixels[:] = data
            image.file_format = 'PNG'
            image.save()
        except Exception as ex:
            print(f'Failed to save preview image: {ex}')
        finally:
            bpy.data.images.remove(image, do_unlink=True)


def extract_assets(blend: str, lib: str, cache_path: str) -> list[dict[str, Any]]:
    """
    Find all node groups, objects, collections and materials which are assets in this file
    and create an exportable entry for each one.
    """
    assets = []

    # Collect all resources in current file.
    all_resources = bpy.data.node_groups[:] + bpy.data.objects[:] + bpy.data.collections[:] + bpy.data.materials[:]

    # Filter all assets.
    all_assets = filter(lambda res: res.asset_data, all_resources)

    # Remove all which are linked.
    all_local = filter(lambda res: not res.library, all_assets)

    for res in all_local:
        # Find type of asset.
        if isinstance(res, bpy.types.Object): asset_type = 'Object'
        elif isinstance(res, bpy.types.Collection): asset_type = 'Collection'
        elif isinstance(res, bpy.types.Material): asset_type = 'Material'
        elif isinstance(res, bpy.types.ShaderNodeTree): asset_type = 'ShaderNodeTree'
        elif isinstance(res, bpy.types.GeometryNodeTree): asset_type = 'GeometryNodeTree'
        else: 
            asset_type = 'Unknown' 

        # Try to extract preview image.
        hash = ''
        if res.preview and res.preview.image_pixels:
            # Build hash from preview image binary data.
            hash_buffer = array.array('i', res.preview.image_pixels[:]).tobytes()
            hash = hashlib.blake2b(hash_buffer).hexdigest()

            # Eventually store the preview externally (if not existing).
            store_preview(
                os.path.join(cache_path, 'previews'),
                hash, 
                res.preview.image_pixels_float[:], 
                res.preview.image_size[:]
            )

        # Get asset info.    
        asset_data = res.asset_data
        assets.append({
            'blend': blend,
            'name': res.name,
            'type': asset_type,
            'library': lib,
            'catalog': asset_data.catalog_id,
            'description': asset_data.description,
            'tags': [ tag.name for tag in res.asset_data.tags ],
            'hash': hash
        })

    return assets


def main(args: list[str]) -> None:
    # Extract script params.
    cache_path = args[0]
    lib_path = args[1]

    # Dir must exist.
    if not os.path.exists(cache_path):
        os.makedirs(cache_path)

    # Unify path of current blend and get last modification time.
    canonical = os.path.normcase(os.path.realpath(os.path.abspath(bpy.data.filepath)))
    last_mod = os.stat(canonical).st_mtime

    # Find asset cache file name and load it.
    print(f'Load asset cache: {canonical}')
    cache_file = os.path.join(cache_path, 'asset_cache_v2.json')
    cache = load_cache(cache_file)

    # Remove entry from current blend.
    new_cache = []
    for lib in cache:
        if lib['name'] != canonical:
            new_cache.append(lib)

    # Extract info of all assets in current blend.
    print(f'Extract asset info from: {canonical}')
    entries = extract_assets(canonical, lib_path, cache_path)

    # Create structure for current blend and append to new cache.
    js = {
        'name': canonical,
        'entry': {
            'last_mod': last_mod,
            'assets': entries
        }
    }
    new_cache.append(js)

    """
    # Hash of source blend.
    hash = hashlib.blake2s(canonical.encode('utf-8')).hexdigest()

    # Path to cache file folder.
    output_path = os.path.join(cache_path, 'asset_cache_v3')
    # Dir must exist.
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # Find file to store the current .blend's asset info.
    output = os.path.join(output_path, hash + '.json')
    """

    print(f'   Store asset cache: {canonical}')
    save_cache(cache_file, { 'libs': new_cache })


if __name__ == "__main__":
    if '--' not in sys.argv:
        argv = [ 'c:/Users/thoma/.asset-wizard-pro/', 'C:\\data\\tmp\\Blender\\tmp\\asset_lib_4.2' ]
    else:
        argv = sys.argv[sys.argv.index('--') + 1:]  # get all args after "--"
    main(argv)
