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

import bpy, os, shutil

from typing import Dict, List, Tuple

from .asset_registry import AssetInfoCache, AssetEntry, asset_preview_image
from .enum_registry import EnumRegistry
from .previews_registry import PreviewsRegistry
from ..utils.io import os_path
from ..utils.common import format_bytes


class ResourceListsRegistry:
    """
    Manages various lists used in the UI.
    Libraries, catalogs, .blend statistics, as well as infos to 
    available nodes in active asset libs.
    """
    def __init__(self):
        PreviewsRegistry.instance().register('assets')

        # Asset preview images.
        self.__previews = {} # type: Dict[str, int]
        # Asset libraries: [ Library, [Blendfile-List (fullpath, name, fullpath)] ]
        self.__library_files = {} # type: Dict[str, List[Tuple[str, str, str]]]

        # Catalog entries: [ Library, [Catalog-List (uuid, visual name, description)]] -> EnumProperty list.
        self.__catalogs = {} # type: Dict[str, List[Tuple[str, str, str]]]

        # Catalog entries: [ Library, [Catalog-List (uuid, path, simple name)]]
        self.__catalogs_info = {} # type: Dict[str, List[Tuple[str, str, str]]]

        # Flat list of catalogs over all libraries -> EnumProperty list.
        self.__catalogs_flat = [] # type: List[Tuple[str, str, str]] 

        # Statistics to show in UI for one .blend in library (number of assets).
        self.__stats = {} # type: Dict[str, List[Tuple[str, str]]]

        # Linear lists of all entries by type.
        self.__shader_nodes = [] # type: List[AssetEntry]
        self.__geometry_nodes = [] # type: List[AssetEntry]
        self.__materials = [] # type: List[AssetEntry]
        self.__col_or_objects = [] # type: List[AssetEntry]

        # Assets by it's uuid.
        self.__all_assets_by_uuid = {} # type: Dict[str, AssetEntry]

        # List of catalogs that contain at least one of the respective asset type.
        # Tuple is usable for EnumProperty (from catalogs_info).
        self.__shader_catalogs = [] # type: List[Tuple[str, str, str]]
        self.__geometry_catalogs = [] # type: List[Tuple[str, str, str]]
        self.__material_catalogs = [] # type: List[Tuple[str, str, str]]
        self.__col_or_object_catalogs = [] # type: List[Tuple[str, str, str]]

        # For each catalog, list of assets valid for EnumProperty, id is uuid from asset.
        # Entries added on first request, than reused.
        self.__shaders_by_catalog = {} # type: Dict[str, List[Tuple[str, str, str]]]
        self.__geometries_by_catalog = {} # type: Dict[str, List[Tuple[str, str, str]]]
        self.__materials_by_catalog = {} # type: Dict[str, List[Tuple[str, str, str, int, int]]]
        self.__col_or_objs_by_catalog = {} # type: Dict[str, List[Tuple[str, str, str, int, int]]]

        # List of categories (first asset tag) that contain at least one of the respective asset type.
        # Tuple is usable for EnumProperty (from catalogs_info).
        self.__shader_categories = [] # type: List[Tuple[str, str, str]]
        self.__geometry_categories = [] # type: List[Tuple[str, str, str]]
        self.__material_categories = [] # type: List[Tuple[str, str, str]]
        self.__col_or_object_categories = [] # type: List[Tuple[str, str, str]]

        # For each category, list of assets valid for EnumProperty, id is uuid from asset.
        # Entries added on first request, than reused.
        self.__shaders_by_category = {} # type: Dict[str, List[Tuple[str, str, str]]]
        self.__geometries_by_category = {} # type: Dict[str, List[Tuple[str, str, str]]]
        self.__materials_by_category = {} # type: Dict[str, List[Tuple[str, str, str, int, int]]]
        self.__col_or_objs_by_category = {} # type: Dict[str, List[Tuple[str, str, str, int, int]]]

        # List of tocs (first letter of name) that contain at least one of the respective asset type.
        # Tuple is usable for EnumProperty (from catalogs_info).
        self.__shader_tocs = [] # type: List[Tuple[str, str, str]]
        self.__geometry_tocs = [] # type: List[Tuple[str, str, str]]
        self.__material_tocs = [] # type: List[Tuple[str, str, str]]
        self.__col_or_object_tocs = [] # type: List[Tuple[str, str, str]]

        # For each toc, list of assets valid for EnumProperty, id is uuid from asset.
        # Entries added on first request, than reused.
        self.__shaders_by_toc = {} # type: Dict[str, List[Tuple[str, str, str]]]
        self.__geometries_by_toc = {} # type: Dict[str, List[Tuple[str, str, str]]]
        self.__materials_by_toc = {} # type: Dict[str, List[Tuple[str, str, str, int, int]]]
        self.__col_or_objs_by_toc = {} # type: Dict[str, List[Tuple[str, str, str, int, int]]]

        self.update(False)


    def __update_library_files(self):
        """
        Build list of all asset .blend files.
        """
        self.__library_files.clear()
        for l in bpy.context.preferences.filepaths.asset_libraries:
            if not os.path.exists(l.path): continue
            path = os_path(l.path)
            r = []
            if l.path:
                for _f in os.listdir(path):
                    if _f.endswith('.blend'):
                        f = os.path.join(path, _f)
                        s = os.path.splitext(_f)[0]
                        r.append((f, s, f))
            self.__library_files[l.path] = sorted(r, key=lambda x: x[1])


    def library_files(self, library_path: str):
        """
        List of all .blend in a single library.
        """
        return self.__library_files.get(library_path, list())


    def __update_catalogs(self):
        """
        Parse catalog files in all repositories.
        """
        self.__catalogs.clear()
        self.__catalogs_info.clear()
        self.__catalogs_flat.clear()

        for lib in bpy.context.preferences.filepaths.asset_libraries:
            if not os.path.exists(lib.path): continue
            catalog_file = os.path.join(os_path(lib.path), 'blender_assets.cats.txt')
            r, i = [], []
            if os.path.exists(catalog_file):
                with open(catalog_file, 'r', encoding='utf-8') as f:
                    for l in f.readlines():
                        l = l.strip()
                        if not l or l[0] == '#' or l.startswith('VERSION'):
                            continue
                        entry = l.split(':')
                        if len(entry) > 2:
                            short_name = (len(entry[1].split('/')) - 1) * '  ' + entry[1].split('/')[-1]
                            long_name = entry[1].replace('/', ' / ')
                            r.append((entry[0], short_name, f'{entry[1]} in {lib.name}'))
                            i.append((entry[0], entry[1], entry[2]))
                            self.__catalogs_flat.append(((entry[0], long_name, f'{long_name} in {lib.name}')))

            r.append(('00000000-0000-0000-0000-000000000000', 'Unassigned', f'Unassigned in {lib.name}'))
            i.append(('00000000-0000-0000-0000-000000000000', 'Unassigned', 'Unassigned'))
            self.__catalogs_flat.append(('00000000-0000-0000-0000-000000000000', 'Unassigned', f'Unassigned in {lib.name}'))

            self.__catalogs[lib.path] = r
            self.__catalogs_info[lib.path] = i

        # Remove double entries by id from catalog_flat.
        unique_dict = {t[0]: t for t in reversed(self.__catalogs_flat)}
        self.__catalogs_flat = sorted(unique_dict.values(), key=lambda x: x[1])



    def catalogs(self, library_path: str):
        """
        All catalogs in a given library.
        """
        return self.__catalogs[library_path]
    

    def catalog_info(self, library_path: str, catalog_uuid: str):
        """
        Get line from catalogs file, split into the 3 components or None.
        """
        for e in self.__catalogs_info[library_path]:
            if e[0] == catalog_uuid:
                return e


    def add_catalog(self, library_path: str, id: str, full_name: str, simple_name: str) -> bool:
        """
        Adds new catalog, returns False if already exists.
        """
        for e in self.__catalogs_info[library_path]:
            if e[1] == full_name or e[2] == simple_name:
                return False

        self.__catalogs_info[library_path].append((id, full_name, simple_name))
        
        catalog_file = os.path.join(os_path(library_path), 'blender_assets.cats.txt')
        if os.path.exists(catalog_file):
            shutil.copy(catalog_file, catalog_file + '.bak')

        with open(catalog_file, 'w', encoding='utf-8') as f:
            f.write('# This is an Asset Catalog Definition file for Blender.\n')
            f.write('#\n')
            f.write('# Empty lines and lines starting with `#` will be ignored.\n')
            f.write('# The first non-ignored line should be the version indicator.\n')
            f.write('# Other lines are of the format "UUID:catalog/path/for/assets:simple catalog name"\n\n')
            f.write('VERSION 1\n\n')
            for u, fn, sn in sorted(self.__catalogs_info[library_path], key=lambda x: x[1]):
                f.write(f'{u}:{fn}:{sn}\n')

        self.update(False)
        return True


    def stats(self, blend_file: str):
        """
        Generate name, value pairs for file information to display in UI, cache entries.
        """
        file = os_path(blend_file)
        if file not in self.__stats:
            r = []
            r.append(('File Size', format_bytes(os.path.getsize(file))))
            with bpy.data.libraries.load(file, link=False, assets_only=True) as (data_from, _):
                if len(data_from.objects): r.append(('Asset Objects', f'{len(data_from.objects)}'))
                if len(data_from.collections): r.append(('Asset Collections', f'{len(data_from.collections)}'))
                if len(data_from.materials): r.append(('Asset Materials', f'{len(data_from.materials)}'))
                if len(data_from.node_groups): r.append(('Asset Node Groups', f'{len(data_from.node_groups)}'))

            self.__stats[file] = r

        return self.__stats[file]
    

    def asset_preview(self, asset: AssetEntry) -> int:
        """
        Returns (eventually load) the icon for a given asset.
        """
        # Already loaded?
        if asset.hash in self.__previews:
            return self.__previews[asset.hash]
        
        # Load image.
        image_file = asset_preview_image(asset.hash)
        id = PreviewsRegistry.instance().collection('assets').load(asset.hash, image_file)

        # Add to dict for next use.
        self.__previews[asset.hash] = id

        # .. and return new id.
        return id      
    

    def __build_filtered_structs(
            self, 
            assets: List[AssetEntry], 
            catalog: List[Tuple[str, str, str]],
            category: List[Tuple[str, str, str]],
            toc: List[Tuple[str, str, str]]
        ):
        """
        Build catalog, category and toc lists.
        """
        category.clear()
        toc.clear()

        # Scan all assets for used entries, create unique lists.
        catalogs, categories, tocs = set(), set(), set()
        for a in assets:
            catalogs.add(a.catalog)
            categories.add(a.category())
            tocs.add(a.toc())

        # Create list of known catalogs from all asset paths, we need to preserve order.
        catalog.clear()
        catalog.extend(self.__catalogs_flat)
        # .. and remove all not in catalogs.
        i = 0
        while i < len(catalog):
            if catalog[i][0] in catalogs:
                i += 1
            else:
                del catalog[i]

        # categories is quite simple.
        for cat in sorted(categories):
            category.append((EnumRegistry.id_for_str(cat), cat, cat))

        # Tocs are not to complex too.
        for t in sorted(tocs):
            toc.append((EnumRegistry.id_for_str(t), f'{t}..', f'Assets starting with: {t}'))


    def __to_uuid(self, asset: AssetEntry, with_preview: bool, id: int) -> Tuple[str, str, str]:
        """
        Turn an asset into an EnumProperty entry.
        """
        if with_preview:
            return (asset.uuid, asset.name, asset.description, self.asset_preview(asset), id)
        else:
            return (asset.uuid, asset.name, asset.description)


    def __by_catalog(self, catalog_uuid: str, assets: List[AssetEntry], with_preview: bool) -> List[Tuple[str, str, str]]:
        """
        Filter assets by given catalog and turn them into EnumProperty entry.
        """
        # index must be start with 0 for initial EnumProperty, so filter first, than build
        lst = filter(lambda a: a.catalog == catalog_uuid, assets)
        return [ self.__to_uuid(a, with_preview, i) for i, a in enumerate(lst) ]
    

    def __by_category(self, category_uuid: str, assets: List[AssetEntry], with_preview: bool) -> List[Tuple[str, str, str]]:
        """
        Filter assets by given category and turn them into EnumProperty entry.
        """
        # index must be start with 0 for initial EnumProperty, so filter first, than build
        cat = EnumRegistry.str_for_id(category_uuid)
        lst = filter(lambda a: a.category() == cat, assets)
        return [ self.__to_uuid(a, with_preview, i) for i, a in enumerate(lst) ]    
    

    def __by_toc(self, toc_uuid: str, assets: List[AssetEntry], with_preview: bool) -> List[Tuple[str, str, str]]:
        """
        Filter assets by given category and turn them into EnumProperty entry.
        """
        t = EnumRegistry.str_for_id(toc_uuid)
        lst = filter(lambda a: a.toc() == t, assets)
        return [ self.__to_uuid(a, with_preview, i) for i, a in enumerate(lst) ]     
    

    def update(self, rescan_from_outdated_file: bool):
        """
        Parses all assets from all .blend files in all asset libraries. 
        Uses the asset cache for faster loading.
        Creates all prefiltered lists used in the UI for faster access.
        """
        self.__update_library_files()
        self.__update_catalogs()
        self.__stats.clear()   

        cache = AssetInfoCache.get()
        cache.re_init()
        for lib, filelist in self.__library_files.items():
            for file in [ e[0] for e in filelist]:
                cache.update_info(lib, file, rescan_from_outdated_file)

        # Hold list pointers.
        self.__all_assets_by_uuid.clear()

        for a in cache.assets():
            self.__all_assets_by_uuid[a.uuid] = a

        self.__shader_nodes.clear()
        self.__geometry_nodes.clear()
        self.__materials.clear()
        self.__col_or_objects.clear()

        self.__shader_nodes.extend(cache.assets_by(lambda a: a.type == 'ShaderNodeTree'))
        self.__geometry_nodes.extend(cache.assets_by(lambda a: a.type == 'GeometryNodeTree'))
        self.__materials.extend(cache.assets_by(lambda a: a.type == 'Material'))
        self.__col_or_objects.extend(cache.assets_by(lambda a: a.type in [ 'Object', 'Collection' ]))

        # Reset all, they are newly filled by first access.
        self.__shaders_by_catalog.clear()
        self.__geometries_by_catalog.clear()
        self.__materials_by_catalog.clear()
        self.__col_or_objs_by_catalog.clear()
        self.__shaders_by_category.clear()
        self.__geometries_by_category.clear()
        self.__materials_by_category.clear()
        self.__col_or_objs_by_category.clear()
        self.__shaders_by_toc.clear()
        self.__geometries_by_toc.clear()
        self.__materials_by_toc.clear()
        self.__col_or_objs_by_toc.clear()

        self.__build_filtered_structs(self.__shader_nodes, self.__shader_catalogs, self.__shader_categories, self.__shader_tocs)
        self.__build_filtered_structs(self.__geometry_nodes, self.__geometry_catalogs, self.__geometry_categories, self.__geometry_tocs)
        self.__build_filtered_structs(self.__materials, self.__material_catalogs, self.__material_categories, self.__material_tocs)
        self.__build_filtered_structs(self.__col_or_objects, self.__col_or_object_catalogs, self.__col_or_object_categories, self.__col_or_object_tocs)

        #self.update_properties()


    def shader_nodes(self): return self.__shader_nodes
    def geometry_nodes(self): return self.__geometry_nodes
    def materials(self): return self.__materials
    def col_or_objects(self): return self.__col_or_objects

    def shader_catalogs(self): return self.__shader_catalogs
    def geometry_catalogs(self): return self.__geometry_catalogs
    def materials_catalogs(self): return self.__material_catalogs
    def col_or_object_catalogs(self): return self.__col_or_object_catalogs

    def shader_categories(self): return self.__shader_categories
    def geometry_categories(self): return self.__geometry_categories
    def materials_categories(self): return self.__material_categories
    def col_or_object_categories(self): return self.__col_or_object_categories

    def shader_tocs(self): return self.__shader_tocs
    def geometry_tocs(self): return self.__geometry_tocs
    def materials_tocs(self): return self.__material_tocs
    def col_or_object_tocs(self): return self.__col_or_object_tocs
    

    def asset_by_uuid(self, uuid: str): return self.__all_assets_by_uuid[uuid]


    def shaders_by_catalog(self, catalog_uuid: str) -> List[Tuple[str, str, str]]:
        if catalog_uuid not in self.__shaders_by_catalog:
            self.__shaders_by_catalog[catalog_uuid] = self.__by_catalog(catalog_uuid, self.__shader_nodes, False)
        return self.__shaders_by_catalog[catalog_uuid]

    def geometries_by_catalog(self, catalog_uuid: str) -> List[Tuple[str, str, str]]:
        if catalog_uuid not in self.__geometries_by_catalog:
            self.__geometries_by_catalog[catalog_uuid] = self.__by_catalog(catalog_uuid, self.__geometry_nodes, False)
        return self.__geometries_by_catalog[catalog_uuid]

    def materials_by_catalog(self, catalog_uuid: str) -> List[Tuple[str, str, str]]:
        if catalog_uuid not in self.__materials_by_catalog:
            self.__materials_by_catalog[catalog_uuid] = self.__by_catalog(catalog_uuid, self.__materials, True)
        return self.__materials_by_catalog[catalog_uuid]

    def col_or_objs_by_catalog(self, catalog_uuid: str) -> List[Tuple[str, str, str]]:
        if catalog_uuid not in self.__col_or_objs_by_catalog:
            self.__col_or_objs_by_catalog[catalog_uuid] = self.__by_catalog(catalog_uuid, self.__col_or_objects, True)
        return self.__col_or_objs_by_catalog[catalog_uuid]


    def shaders_by_category(self, category_uuid: str) -> List[Tuple[str, str, str]]:
        if category_uuid not in self.__shaders_by_category:
            self.__shaders_by_category[category_uuid] = self.__by_category(category_uuid, self.__shader_nodes, False)
        return self.__shaders_by_category[category_uuid]

    def geometries_by_category(self, category_uuid: str) -> List[Tuple[str, str, str]]:
        if category_uuid not in self.__geometries_by_category:
            self.__geometries_by_category[category_uuid] = self.__by_category(category_uuid, self.__geometry_nodes, False)
        return self.__geometries_by_category[category_uuid]

    def materials_by_category(self, category_uuid: str) -> List[Tuple[str, str, str]]:
        if category_uuid not in self.__materials_by_category:
            self.__materials_by_category[category_uuid] = self.__by_category(category_uuid, self.__materials, True)
        return self.__materials_by_category[category_uuid]

    def col_or_objs_by_category(self, category_uuid: str) -> List[Tuple[str, str, str]]:
        if category_uuid not in self.__col_or_objs_by_category:
            self.__col_or_objs_by_category[category_uuid] = self.__by_category(category_uuid, self.__col_or_objects, True)
        return self.__col_or_objs_by_category[category_uuid]


    def shaders_by_toc(self, toc_uuid: str) -> List[Tuple[str, str, str]]:
        if toc_uuid not in self.__shaders_by_toc:
            self.__shaders_by_toc[toc_uuid] = self.__by_toc(toc_uuid, self.__shader_nodes, False)
        return self.__shaders_by_toc[toc_uuid]
    
    def geometries_by_toc(self, toc_uuid: str) -> List[Tuple[str, str, str]]:
        if toc_uuid not in self.__geometries_by_toc:
            self.__geometries_by_toc[toc_uuid] = self.__by_toc(toc_uuid, self.__geometry_nodes, False)
        return self.__geometries_by_toc[toc_uuid]

    def materials_by_toc(self, toc_uuid: str) -> List[Tuple[str, str, str]]:
        if toc_uuid not in self.__materials_by_toc:
            self.__materials_by_toc[toc_uuid] = self.__by_toc(toc_uuid, self.__materials, False)
        return self.__materials_by_toc[toc_uuid]

    def col_or_objs_by_toc(self, toc_uuid: str) -> List[Tuple[str, str, str]]:
        if toc_uuid not in self.__col_or_objs_by_toc:
            self.__col_or_objs_by_toc[toc_uuid] = self.__by_toc(toc_uuid, self.__col_or_objects, False)
        return self.__col_or_objs_by_toc[toc_uuid]            


    @staticmethod
    def get() -> 'ResourceListsRegistry':
        global _resource_lists_instance
        return _resource_lists_instance


_resource_lists_instance = ResourceListsRegistry()        
