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

import os
from typing import Any, Dict, List, Tuple, Union

from ..constants import config_file
from .textures.types import ( RepositoryType, TEXTURE_PATH_PBR, TEXTURE_PATH_EXTREME_PBR, 
            TEXTURE_PATH_GRUNGE, TEXTURE_PATH_BRUSH_GRAY, TEXTURE_PATH_IMAGE, texture_path_style_name )
from .textures.texture_entry import TextureEntry
from .textures.texture_repository import TextureRepository
from .textures.image_hash_repository import HashRepository
from .textures.image_lookup import ImageLookup
from .textures.scanners.pbr_directory import PBRDirectory
from .textures.scanners.gray_directory import GrayDirectory
from .textures.scanners.brush_directory import BrushDirectory
from .textures.scanners.image_directory import ImageDirectory
from .textures.scanners.pbr_scanner import PBRLinearScanner
from .textures.scanners.extreme_pbr_scanner import PBRExtremeScanner
from .textures.scanners.gray_scanner import GrayLinearScanner
from .textures.scanners.brush_scanner import BrushLinearScanner
from .textures.scanners.image_scanner import ImageLinearScanner
from ..utils.tools import Measure
from ..utils.dev import dbg, inf, log_exception
from ..utils.io import read_json, unified_path, write_json, decode_json


class TextureRegistry:
    """
    The root of all texture information, hosts all texture
    repositories from preferences.
    """
    def init(self): 
        # Image trees.
        self.__repositories = {} # type: Dict[RepositoryType, Dict[str, TextureRepository]]
        self.__repositories[RepositoryType.Gray] = {} # type: Dict[str, TextureRepository]
        self.__repositories[RepositoryType.PBR] = {} # type: Dict[str, TextureRepository]
        self.__repositories[RepositoryType.Brush] = {} # type: Dict[str, TextureRepository]
        self.__repositories[RepositoryType.Image] = {} # type: Dict[str, TextureRepository]

        # Enums for UI.
        self.__enums = {} # type: Dict[RepositoryType, List[Any]]
        self.__enums[RepositoryType.Gray] = []
        self.__enums[RepositoryType.PBR] = []
        self.__enums[RepositoryType.Brush] = []
        self.__enums[RepositoryType.Image] = []

        self.__hash_repository = HashRepository()
        self.__image_lookup = ImageLookup()

        # Load all information from previous scans. We load this into
        # a separate dict first and in scan_repository(), check here if we have the
        # information. In this case, copy the respective path's data into self.repositories.
        # Loading the cache NOT into self.repositories directly is simple: 
        # If the cache contains a repositority, that is actually not used
        # anymore, it still resides in cache. But it shoudln't be self.repositories.
        # The cache file is newly created in scan_repositories, so when closing Blender
        # it is finally up2date.
        self.__cached_repositories = self.load()

        dbg('TextureRegistry created')


    def dispose(self):
        """
        Free all resources.
        """
        self.__image_lookup.clear()

        for repo in self.__repositories.values():
            for r in repo.values():
                r.dispose(self.__image_lookup)

        self.__repositories[RepositoryType.Gray].clear()
        self.__repositories[RepositoryType.PBR].clear()
        self.__repositories[RepositoryType.Brush].clear()
        self.__repositories[RepositoryType.Image].clear()
        self.__enums[RepositoryType.Gray].clear()
        self.__enums[RepositoryType.PBR].clear() 
        self.__enums[RepositoryType.Brush].clear()
        self.__enums[RepositoryType.Image].clear()

        dbg('TextureRegistry disposed')


    def __update_repository_lists(self, type: RepositoryType):
        """
        Rebuild category enum list.
        """
        # Collect categories from all repositories of this type.
        flat_categories = [] 
        for r in self.__repositories[type].values():
            flat_categories.extend(r.get_category_enums())

        # Sort by description and store in out list.
        self.__enums[type].clear()
        self.__enums[type].extend(sorted(flat_categories, key=lambda x: x[2]))
        
        # Update image lookup as well.
        self.__image_lookup.clear()
        for repo in self.__repositories.values():
            for r in repo.values():
                r.update_image_lookup(self.__image_lookup)


    def __add_pbr_repository(self, base_path: str, prefix: str, style: int, categories: List[PBRDirectory]):
        new_repo = TextureRepository(base_path, RepositoryType.PBR, prefix, style)
        new_repo.add_pbr(categories, self.__hash_repository)
        self.__repositories[RepositoryType.PBR][base_path] = new_repo
        self.__update_repository_lists(RepositoryType.PBR)


    def __add_gray_repository(self, base_path: str, prefix: str, style: int, categories: List[GrayDirectory]):
        new_repo = TextureRepository(base_path, RepositoryType.Gray, prefix, style)
        new_repo.add_gray(categories, self.__hash_repository)
        self.__repositories[RepositoryType.Gray][base_path] = new_repo
        self.__update_repository_lists(RepositoryType.Gray)        


    def __add_brush_repository(self, base_path: str, prefix: str, style: int, categories: List[BrushDirectory]):
        new_repo = TextureRepository(base_path, RepositoryType.Brush, prefix, style)
        new_repo.add_brush(categories, self.__hash_repository)
        self.__repositories[RepositoryType.Brush][base_path] = new_repo
        self.__update_repository_lists(RepositoryType.Brush)        


    def __add_image_repository(self, base_path: str, prefix: str, style: int, categories: List[ImageDirectory]):
        new_repo = TextureRepository(base_path, RepositoryType.Image, prefix, style)
        new_repo.add_image(categories, self.__hash_repository)
        self.__repositories[RepositoryType.Image][base_path] = new_repo
        self.__update_repository_lists(RepositoryType.Image)


    def scan_repository(self, path: str, prefix: str, style: int, force_scan: bool = False):
        # Bring path into a unique scheme.
        path = unified_path(path)

        # Convert style to repository type.
        type = {
            TEXTURE_PATH_GRUNGE: RepositoryType.Gray,
            TEXTURE_PATH_PBR: RepositoryType.PBR,
            TEXTURE_PATH_EXTREME_PBR: RepositoryType.PBR,
            TEXTURE_PATH_BRUSH_GRAY: RepositoryType.Brush,
            TEXTURE_PATH_IMAGE: RepositoryType.Image
        }[style]

        # If the path is already there (from load()) ...
        if path in self.__repositories[type]:
            # This is true if your starts manual rescan.
            if force_scan:
                inf(f'Rescan {path} due to force')
                del self.__repositories[type][path]
            else:
                inf(f'Skip scanning {path}, already loaded')
                self.__update_repository_lists(type)
                return
        else:
            # We might have the information from a previous run.
            if not force_scan and type in self.__cached_repositories and path in self.__cached_repositories[type]:
                inf(f'Skip scanning {path}, found in cache')
                self.__repositories[type][path] = self.__cached_repositories[type][path]
                self.__update_repository_lists(type)
                return


        inf(f'Start scanning {path}')
        with Measure(f'Scan textures in {path} (style={texture_path_style_name(style)})'):
            if style == TEXTURE_PATH_GRUNGE:
                self.__add_gray_repository(path, prefix, style, GrayLinearScanner.scan(path))
            elif style == TEXTURE_PATH_PBR:
                self.__add_pbr_repository(path, prefix, style, PBRLinearScanner.scan(path))
            elif style == TEXTURE_PATH_EXTREME_PBR:
                self.__add_pbr_repository(path, prefix, style, PBRExtremeScanner.scan(path))
            elif style == TEXTURE_PATH_BRUSH_GRAY:
                self.__add_brush_repository(path, prefix, style, BrushLinearScanner.scan(path))
            elif style == TEXTURE_PATH_IMAGE:
                self.__add_image_repository(path, prefix, style, ImageLinearScanner.scan(path))

        self.__hash_repository.save()
        self.save()


    def remove_repository(self, path: str):
        """
        Remove a path from internal list.
        """
        e = self.__get_repository_by_path(path)
        if e:
            type, fixed_path, repo = e
            repo.dispose(self.__image_lookup)
            del self.__repositories[type][fixed_path]
            self.__update_repository_lists(type)


    def rescan_repository(self, path: str):
        """
        Rescan textures from a repository, by removing a readding it.
        """
        e = self.__get_repository_by_path(path)
        if e:
            _, fixed_path, repo = e
            self.remove_repository(fixed_path)
            self.scan_repository(path, repo.prefix, repo.style, True)


    def update_repository_prefix(self, path: str, prefix: str):
        """
        Called when prefix has been changed.
        """
        e = self.__get_repository_by_path(path)
        if e:
            type, _, repo = e
            repo.update_prefix(prefix)
            self.__update_repository_lists(type)


    def __get_repository_by_path(self, path: str) -> Tuple[RepositoryType, str, TextureRepository]:
        """
        Return type and repository by path.
        """
        for type, repos in self.__repositories.items():
            for p in repos.keys():
                if os.path.samefile(path, p):
                    return (type, p, repos[p])


    def __get_repository_for_category(self, type: RepositoryType, category: str) -> Union[Tuple[TextureRepository, str], None]:
        """
        Returns matching repo for given info. 
        Category is a key, refering TextureRepository.__key.
        """
        try:
            for repo in self.__repositories[type].values():
                if repo.has_category(category):
                    return (repo, category)
        except Exception as e:
            log_exception(e, context_msg="Failed to get repository for category")
        return None


    def create_update_previews(self, type: RepositoryType, category: str, missing_only: bool, size: int, whole_repository: bool = False):
        """
        Re-renders previews for all images in category/repo.
        """
        repo, category = self.__get_repository_for_category(type, category)
        return repo.create_update_previews(category, missing_only, size, whole_repository)


    def get_category_enums(self, type: RepositoryType) -> List[Tuple[str, str, str]]:
        """
        Get enum list for categories to use in EnumProperty.
        """
        return self.__enums[type]


    def get_image_enums(self, type: RepositoryType, category_key: str) -> List[Tuple[str, str, str, int, int]]:
        """
        Get image enum list for category.
        """
        repo, category_key = self.__get_repository_for_category(type, category_key)
        return repo.get_image_enums(category_key)


    def get_image_list(self, type: RepositoryType, category_key: str, with_preview: bool = False, with_gpu: bool = False) -> List[TextureEntry]:
        """
        Get image enum list for category.
        """
        repo, category_key = self.__get_repository_for_category(type, category_key)
        return repo.get_entries(category_key, with_preview, with_gpu)      


    def get_image_by_snw_info(self, snw_info: str, with_preview: bool = False, with_gpu: bool = False) -> TextureEntry:
        """
        Decode info from snw_info (json) and return entry if exists.
        """
        js = decode_json(snw_info)
        repository_type = RepositoryType.PBR if js.get('type', 'P') == 'P' else RepositoryType.Gray
        category = js.get('category', None)
        image = js.get('pbr', None) if repository_type == RepositoryType.PBR else js.get('grunge', None)
        if category and image:
            return self.get_image(repository_type, category, image, with_preview, with_gpu)


    def get_image(self, type: RepositoryType, category_id: str, image: str, with_preview: bool = False, with_gpu: bool = False) -> TextureEntry:
        """
        Return the image found using the described data.
        """
        info = self.__get_repository_for_category(type, category_id)
        if info:
            repo, category_id = info
            return repo.get_entry(category_id, image, with_preview, with_gpu)


    def get_by_bin_hash(self, bin_hash: str) -> Tuple[Union[str, None], Union[str, None]]:
        """
        Find entry in image cache, returns category key and image key.
        """
        return self.__image_lookup.get(bin_hash)


    def load(self) -> Dict[RepositoryType, Dict[str, TextureRepository]]:
        """
        Called on startup, this load all repository information
        from cache-file, to speed up Blender startup.
        """
        cache_file = config_file('repo_cache.json')
        inf(f'Load texture info from {cache_file}')
        data = read_json(cache_file)

        cached_repositories = {} # type: Dict[RepositoryType, Dict[str, TextureRepository]]

        # Only load this exact version, otherwise a new one is created
        # from automatic scanning.
        if data.get('version') == 1:
            for repository_by_type in data.get('repositories', []):
                type = repository_by_type.get('type', 0)
                entries = repository_by_type.get('entries', [])
                repositories_by_path = {}
                for e in entries:
                    r = TextureRepository.load(e)
                    repositories_by_path[r.name] = r

                if type not in cached_repositories:
                    cached_repositories[type] = {} # type: Dict[str, TextureRepository]
                for k, v in repositories_by_path.items():
                    cached_repositories[type][k] = v

        return cached_repositories


    def save(self):
        """
        Store repository into a file, so we dont have to scan on each startup (which
        may take very long). This is loaded in __init__(). Single entries can be
        forced to reload, when scan_repository is called with force_scan=True.
        """
        data = {} # type: Dict[str, Any]

        # Add a version number. If internal structure changes, we increase
        # the version number, so loading is skipped, and all dirs are scanned
        # automatically.
        data['version'] = 1

        repository_by_type = [] # type: List[Dict]
        for k, v in self.__repositories.items():
            repository_by_path = [] # type: List[Dict]
            for _, pv in v.items():
                repository_by_path.append(pv.save())
            repository_by_type.append({
                'type': k,
                'entries': repository_by_path
            })

        data['repositories'] = repository_by_type

        cache_file = config_file('repo_cache.json')
        inf(f'Save texture info to {cache_file}')
        write_json(config_file('repo_cache.json'), data)


    @staticmethod
    def instance() -> 'TextureRegistry':
        """
        Access as singleton.
        """
        global _instance
        return _instance


_instance = TextureRegistry()
