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

import hashlib, os, json

from typing import Any, Dict, List

from ...utils.io import run_blender
from ...utils.dev import inf
from .types import TextureType
from .image_hash_repository import HashRepository
from .image_lookup import ImageLookup
from .scanners.pbr_entry import PBREntry
from .scanners.gray_entry import GrayEntry
from .scanners.brush_entry import BrushEntry
from .scanners.image_entry import ImageEntry


class TextureInfo:
    def __init__(
        self, 
        hash_repository: HashRepository = None, 
        pbr: PBREntry = None, 
        gray: GrayEntry = None, 
        brush: BrushEntry = None,
        image: ImageEntry = None,
        bin_hash: List[str] = None,
        bin_hash_v2: List[str] = None
        ):
        self.pbr, self.gray, self.brush, self.image = pbr, gray, brush, image

        self.bin_hash = bin_hash if bin_hash else []
        self.bin_hash_v2 = bin_hash_v2 if bin_hash_v2 else []

        if hash_repository:
            # Create hashes for fast lookup.
            if self.pbr: 
                self.__add_hash(hash_repository, self.pbr.diffuse)
                for v in self.pbr.variants:
                    self.__add_hash(hash_repository, v.diffuse)
            if self.gray: 
                # Gray files may be built using channels, all should have a different cache.
                self.__add_hash(hash_repository, self.gray.image, f'@{self.gray.channel}')
            if self.brush: 
                self.__add_hash(hash_repository, self.brush.image)
            if self.image: 
                self.__add_hash(hash_repository, self.image.image)


    def __add_hash(self, hash_repository: HashRepository, hash_file: str, hash_ext: str = ''):
        # Get it from cache or build it (and store it to cache).
        # Check cache first.
        bin_hash = hash_repository.get_file_hash(hash_file)
        if not bin_hash:
            # Not in cache, build it.
            inf(f'Create hash for {hash_file}')
            with open(hash_file, 'rb') as f:
                hash_data = f.read(16384)
                bin_hash = hashlib.sha256(hash_data).hexdigest()
            # Store in cache.
            hash_repository.add_file_hash(hash_file, bin_hash)
        bin_hash += hash_ext 
        self.bin_hash.append(bin_hash)


        # Do the same for v2.
        bin_hash_v2 = hash_repository.get_file_hash_v2(hash_file)
        if not bin_hash_v2:
            # Not in cache, build it.
            inf(f'Create hash v2 for {hash_file}')
            with open(hash_file, 'rb') as f:
                hash_data = f.read()
                bin_hash_v2 = hashlib.sha256(hash_data).hexdigest()
            # Store in cache.
            hash_repository.add_file_hash_v2(hash_file, bin_hash_v2)
        bin_hash_v2 += hash_ext 
        self.bin_hash_v2.append(bin_hash_v2)

        return (bin_hash, bin_hash_v2)


    def dispose(self, image_lookup: ImageLookup):
        for e in self.bin_hash: image_lookup.remove_entry(e)
        for e in self.bin_hash_v2: image_lookup.remove_entry_v2(e)


    def type(self) -> TextureType:
        if self.pbr: return TextureType.PBRSet
        if self.brush: return TextureType.Brush
        if self.image: return TextureType.Image
        if self.gray.channel == -1: return TextureType.Gray1
        return TextureType.Gray4


    def name(self) -> str:
        if self.pbr: return self.pbr.name
        if self.gray: return self.gray.name
        if self.brush: return self.brush.name
        if self.image: return self.image.name
        return ""


    def get_bin_hash(self, index: int = 0) -> str: return self.bin_hash_v2[index]


    def create_update_preview(self, missing_only: bool, size: int) -> bool:
        """
        Eventually re-render the preview. Returns true if image has been updated.
        """
        preview = ''
        if self.pbr: preview = self.pbr.preview
        if self.gray: preview = self.gray.preview
        if self.brush: preview = self.brush.preview
        if self.image: preview = self.image.preview

        if not missing_only or not os.path.exists(preview):
            if self.pbr:
                script = os.path.join(os.path.dirname(__file__), '..', '..', 'preview', 'render_preview_pbr.py')
                blend = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'preview', 'preview-pbr.blend')
                data = json.dumps({
                    'diffuse': self.pbr.diffuse,
                    'metal': self.pbr.metal,
                    'roughness': self.pbr.roughness,
                    'normal': self.pbr.normal,
                    'height': self.pbr.height
                })
            elif self.brush:
                script = os.path.join(os.path.dirname(__file__), '..', '..', 'preview', 'render_preview_brush.py')
                blend = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'preview', 'preview-brush.blend')
                data = json.dumps({
                    'image': self.brush.image,
                })
            elif self.image:
                script = os.path.join(os.path.dirname(__file__), '..', '..', 'preview', 'render_preview_image.py')
                blend = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'preview', 'preview-image.blend')
                data = json.dumps({
                    'image': self.image.image,
                })
            else:
                script = os.path.join(os.path.dirname(__file__), '..', '..', 'preview', 'render_preview_gray.py')
                blend = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'preview', 'preview-gray.blend')
                data = json.dumps({
                    'image': self.gray.image,
                    'channel': self.gray.channel
                })
            
            args = [
                '--background',
                '--factory-startup',
                blend,
                '--python',
                script,
                '--',
                data,
                preview,
                str(size)
            ]
            run_blender(args)
            return True

        return False


    def valid_preview_image(self) -> str:
        """
        Return either the preview image or the original (diffuse),
        so a thumb can be generated.
        """
        if self.pbr:
            p = self.pbr.preview 
            a = self.pbr.diffuse
        elif self.gray:
            p = self.gray.preview
            a = self.gray.image
        elif self.image:
            p = self.image.preview
            a = self.image.image
        else:
            p = self.brush.preview
            a = self.brush.image
        return p if p and os.path.exists(p) else a


    @staticmethod
    def load(data: Dict[str, Any]) -> 'TextureInfo':
        gray, pbr, brush, image = None, None, None, None
        bin_hash, bin_hash_v2 = [], []
        if 'gray' in data: gray = GrayEntry.load(data['gray'])
        if 'pbr' in data: pbr = PBREntry.load(data['pbr'])
        if 'brush' in data: brush = BrushEntry.load(data['brush'])
        if 'image' in data: image = ImageEntry.load(data['image'])
        if 'hash' in data: bin_hash = data['hash']
        if 'hash_v2' in data: bin_hash_v2 = data['hash_v2']
        return TextureInfo(
            gray=gray,
            pbr=pbr,
            brush=brush,
            image=image,
            bin_hash=bin_hash,
            bin_hash_v2=bin_hash_v2
        )


    def save(self) -> Dict[str, Any]:
        r = {
            'hash': self.bin_hash,
            'hash_v2': self.bin_hash_v2
        }

        if self.gray: r['gray'] = self.gray.save()
        if self.pbr: r['pbr'] = self.pbr.save()
        if self.brush: r['brush'] = self.brush.save()
        if self.image: r['image'] = self.image.save()

        return r
