from hydrus.core.files.images import HydrusImageInit # right up top

import cv2
import hashlib
import io
import numpy
import typing
import warnings

from PIL import ImageFile as PILImageFile
from PIL import Image as PILImage

try:
    
    from pillow_heif import register_heif_opener
    from pillow_heif import register_avif_opener
    
    register_heif_opener(thumbnails=False)
    register_avif_opener(thumbnails=False)
    
    HEIF_OK = True
    
except:
    
    HEIF_OK = False
    

from hydrus.core import HydrusConstants as HC
from hydrus.core import HydrusData
from hydrus.core import HydrusExceptions
from hydrus.core import HydrusGlobals as HG
from hydrus.core.files import HydrusKritaHandling
from hydrus.core.files import HydrusPSDHandling
from hydrus.core.files.images import HydrusImageColours
from hydrus.core.files.images import HydrusImageMetadata
from hydrus.core.files.images import HydrusImageNormalisation
from hydrus.core.files.images import HydrusImageOpening

def EnableLoadTruncatedImages():
    
    if hasattr( PILImageFile, 'LOAD_TRUNCATED_IMAGES' ):
        
        # this can now cause load hangs due to the trunc load code adding infinite fake EOFs to the file stream, wew lad
        # hence debug only
        PILImageFile.LOAD_TRUNCATED_IMAGES = True
        
        return True
        
    else:
        
        return False
        
    

OLD_PIL_MAX_IMAGE_PIXELS = PILImage.MAX_IMAGE_PIXELS
PILImage.MAX_IMAGE_PIXELS = None # this turns off decomp check entirely, wew

if cv2.__version__.startswith( '2' ):
    
    CV_IMREAD_FLAGS_PNG = cv2.CV_LOAD_IMAGE_UNCHANGED
    CV_IMREAD_FLAGS_JPEG = CV_IMREAD_FLAGS_PNG
    CV_IMREAD_FLAGS_WEIRD = CV_IMREAD_FLAGS_PNG
    
    CV_JPEG_THUMBNAIL_ENCODE_PARAMS = []
    CV_PNG_THUMBNAIL_ENCODE_PARAMS = []
    
else:
    
    # allows alpha channel
    CV_IMREAD_FLAGS_PNG = cv2.IMREAD_UNCHANGED
    # this preserves colour info but does EXIF reorientation and flipping
    CV_IMREAD_FLAGS_JPEG = cv2.IMREAD_ANYDEPTH | cv2.IMREAD_ANYCOLOR
    # this seems to allow weirdass tiffs to load as non greyscale, although the LAB conversion 'whitepoint' or whatever can be wrong
    CV_IMREAD_FLAGS_WEIRD = CV_IMREAD_FLAGS_PNG
    
    CV_JPEG_THUMBNAIL_ENCODE_PARAMS = [ cv2.IMWRITE_JPEG_QUALITY, 92 ]
    CV_PNG_THUMBNAIL_ENCODE_PARAMS = [ cv2.IMWRITE_PNG_COMPRESSION, 9 ]
    

PIL_ONLY_MIMETYPES = { HC.ANIMATION_GIF, HC.IMAGE_ICON, HC.IMAGE_WEBP, HC.IMAGE_QOI, HC.IMAGE_BMP }.union( HC.PIL_HEIF_MIMES )

def MakeClipRectFit( image_resolution, clip_rect ):
    
    ( im_width, im_height ) = image_resolution
    ( x, y, clip_width, clip_height ) = clip_rect
    
    x = max( 0, x )
    y = max( 0, y )
    
    clip_width = min( clip_width, im_width )
    clip_height = min( clip_height, im_height )
    
    if x + clip_width > im_width:
        
        x = im_width - clip_width
        
    
    if y + clip_height > im_height:
        
        y = im_height - clip_height
        
    
    return ( x, y, clip_width, clip_height )
    
def ClipNumPyImage( numpy_image: numpy.array, clip_rect ):
    
    if len( numpy_image.shape ) == 3:
        
        ( im_height, im_width, depth ) = numpy_image.shape
        
    else:
        
        ( im_height, im_width ) = numpy_image.shape
        
    
    ( x, y, clip_width, clip_height ) = MakeClipRectFit( ( im_width, im_height ), clip_rect )
    
    return numpy_image[ y : y + clip_height, x : x + clip_width ]
    

def ClipPILImage( pil_image: PILImage.Image, clip_rect ):
    
    ( x, y, clip_width, clip_height ) = MakeClipRectFit( pil_image.size, clip_rect )
    
    return pil_image.crop( box = ( x, y, x + clip_width, y + clip_height ) )
    

def GenerateNumPyImage( path, mime, force_pil = False ) -> numpy.array:
    
    if HG.media_load_report_mode:
        
        HydrusData.ShowText( 'Loading media: ' + path )
        
    
    if mime == HC.APPLICATION_PSD:
        
        if HG.media_load_report_mode:
            
            HydrusData.ShowText( 'Loading PSD' )
            
        
        pil_image = HydrusPSDHandling.MergedPILImageFromPSD( path )
        
        return GenerateNumPyImageFromPILImage( pil_image )
        
    
    if mime == HC.APPLICATION_KRITA:
        
        if HG.media_load_report_mode:
            
            HydrusData.ShowText( 'Loading KRA' )
            
        
        pil_image = HydrusKritaHandling.MergedPILImageFromKra( path )
        
        return GenerateNumPyImageFromPILImage( pil_image )
        
    
    if mime in PIL_ONLY_MIMETYPES:
        
        force_pil = True
        
    
    if not force_pil:
        
        pil_image = HydrusImageOpening.RawOpenPILImage( path )
        
        if pil_image.mode == 'LAB':
            
            force_pil = True
            
        
        if HydrusImageMetadata.HasICCProfile( pil_image ):
            
            if HG.media_load_report_mode:
                
                HydrusData.ShowText( 'Image has ICC, so switching to PIL' )
                
            
            force_pil = True
            
        
    
    if force_pil:
        
        if HG.media_load_report_mode:
            
            HydrusData.ShowText( 'Loading with PIL' )
            
        
        pil_image = GeneratePILImage( path )
        
        numpy_image = GenerateNumPyImageFromPILImage( pil_image )
        
    else:
        
        if HG.media_load_report_mode:
            
            HydrusData.ShowText( 'Loading with OpenCV' )
            
        
        if mime in ( HC.IMAGE_JPEG, HC.IMAGE_TIFF ):
            
            flags = CV_IMREAD_FLAGS_JPEG
            
        elif mime == HC.IMAGE_PNG:
            
            flags = CV_IMREAD_FLAGS_PNG
            
        else:
            
            flags = CV_IMREAD_FLAGS_WEIRD
            
        
        numpy_image = cv2.imread( path, flags = flags )
        
        if numpy_image is None: # doesn't support some random stuff
            
            if HG.media_load_report_mode:
                
                HydrusData.ShowText( 'OpenCV Failed, loading with PIL' )
                
            
            pil_image = GeneratePILImage( path )
            
            numpy_image = GenerateNumPyImageFromPILImage( pil_image )
            
        else:
            
            numpy_image = HydrusImageNormalisation.DequantizeFreshlyLoadedNumPyImage( numpy_image )
            
            numpy_image = HydrusImageNormalisation.StripOutAnyUselessAlphaChannel( numpy_image )
            
        
    
    return numpy_image
    
def GenerateNumPyImageFromPILImage( pil_image: PILImage.Image, strip_useless_alpha = True ) -> numpy.array:
    
    try:
        
        # this seems to magically work, I guess asarray either has a match for Image or Image provides some common shape/datatype properties that it can hook into
        numpy_image = numpy.asarray( pil_image )
        
    except IOError:
        
        raise HydrusExceptions.DamagedOrUnusualFileException( 'Looks like a truncated file that PIL could not handle!' )
        
    
    if numpy_image.shape == ():
        
        raise HydrusExceptions.DamagedOrUnusualFileException( 'Looks like a weird truncated file!' )
        
    
    if strip_useless_alpha:
        
        numpy_image = HydrusImageNormalisation.StripOutAnyUselessAlphaChannel( numpy_image )
        
    
    return numpy_image
    

def GeneratePILImage( path: typing.Union[ str, typing.BinaryIO ], dequantize = True ) -> PILImage.Image:
    
    pil_image = HydrusImageOpening.RawOpenPILImage( path )
    
    try:
        
        pil_image = HydrusImageNormalisation.RotateEXIFPILImage( pil_image )
        
        if dequantize:
            
            if pil_image.mode in ( 'I', 'F' ):
                
                # 'I' = greyscale, uint16
                # 'F' = float, np.float32
                
                # calling "pil_image.convert( 'L' )" and similar on an I doesn't seem to normalise the intensity, it just blows out the image crazy???
                # so we'll hack it ourselves. these are so rare it is fine if we are a bit weird and lose the extra metadata
                
                numpy_image = GenerateNumPyImageFromPILImage( pil_image )
                
                numpy_image = HydrusImageNormalisation.NormaliseNumPyImageToUInt8( numpy_image )
                
                pil_image = GeneratePILImageFromNumPyImage( numpy_image )
                
            
            # note this destroys animated gifs atm, it collapses down to one frame
            pil_image = HydrusImageNormalisation.DequantizePILImage( pil_image )
            
        
        return pil_image
        
    except IOError:
        
        raise HydrusExceptions.DamagedOrUnusualFileException( 'Looks like a truncated file that PIL could not handle!' )
        
    

def GeneratePILImageFromNumPyImage( numpy_image: numpy.array ) -> PILImage.Image:
    
    if len( numpy_image.shape ) == 2:
        
        ( h, w ) = numpy_image.shape
        
        format = 'L'
        
    else:
        
        ( h, w, depth ) = numpy_image.shape
        
        if depth == 1:
            
            format = 'L'
            
        elif depth == 2:
            
            format = 'LA'
            
        elif depth == 3:
            
            format = 'RGB'
            
        elif depth == 4:
            
            format = 'RGBA'
            
        
    
    pil_image = PILImage.frombytes( format, ( w, h ), numpy_image.data.tobytes() )
    
    return pil_image
    

def GenerateThumbnailNumPyFromStaticImagePath( path, target_resolution, mime ):
    
    numpy_image = GenerateNumPyImage( path, mime )
    
    thumbnail_numpy_image = ResizeNumPyImage( numpy_image, target_resolution )
    
    return thumbnail_numpy_image
    

def GenerateThumbnailBytesFromNumPy( numpy_image ) -> bytes:
    
    if len( numpy_image.shape ) == 2:
        
        depth = 3
        
        convert = cv2.COLOR_GRAY2RGB
        
    else:
        
        ( im_height, im_width, depth ) = numpy_image.shape
        
        numpy_image = HydrusImageNormalisation.StripOutAnyUselessAlphaChannel( numpy_image )
        
        if depth == 4:
            
            convert = cv2.COLOR_RGBA2BGRA
            
        else:
            
            convert = cv2.COLOR_RGB2BGR
            
        
    
    numpy_image = cv2.cvtColor( numpy_image, convert )
    
    ( im_height, im_width, depth ) = numpy_image.shape
    
    if depth == 4:
        
        ext = '.png'
        
        params = CV_PNG_THUMBNAIL_ENCODE_PARAMS
        
    else:
        
        ext = '.jpg'
        
        params = CV_JPEG_THUMBNAIL_ENCODE_PARAMS
        
    
    ( result_success, result_byte_array ) = cv2.imencode( ext, numpy_image, params )
    
    if result_success:
        
        thumbnail_bytes = result_byte_array.tostring()
        
        return thumbnail_bytes
        
    else:
        
        raise HydrusExceptions.CantRenderWithCVException( 'Thumb failed to encode!' )
        
    

def GenerateThumbnailBytesFromPIL( pil_image: PILImage.Image ) -> bytes:
    
    f = io.BytesIO()
    
    if HydrusImageColours.PILImageHasTransparency( pil_image ):
        
        pil_image.save( f, 'PNG' )
        
    else:
        
        pil_image.save( f, 'JPEG', quality = 92 )
        
    
    f.seek( 0 )
    
    thumbnail_bytes = f.read()
    
    f.close()
    
    return thumbnail_bytes
    

def GeneratePNGBytesNumPy( numpy_image ) -> bytes:
    
    ( im_height, im_width, depth ) = numpy_image.shape

    ext = '.png'

    if depth == 4:
        
        convert = cv2.COLOR_RGBA2BGRA
        
    else:
        
        convert = cv2.COLOR_RGB2BGR
        
    
    numpy_image = cv2.cvtColor( numpy_image, convert )
    
    ( result_success, result_byte_array ) = cv2.imencode( ext, numpy_image )
    
    if result_success:
        
        return result_byte_array.tostring()
        
    else:
        
        raise HydrusExceptions.CantRenderWithCVException( 'Image failed to encode!' )
        
    

def GetImagePixelHash( path, mime ) -> bytes:
    
    numpy_image = GenerateNumPyImage( path, mime )
    
    return GetImagePixelHashNumPy( numpy_image )
    

def GetImagePixelHashNumPy( numpy_image ):
    
    return hashlib.sha256( numpy_image.data.tobytes() ).digest()
    

def GetImageResolution( path, mime ):
    
    # PIL first here, rather than numpy, as it loads image headers real quick
    try:
        
        pil_image = GeneratePILImage( path, dequantize = False )
        
        ( width, height ) = pil_image.size
        
    except HydrusExceptions.DamagedOrUnusualFileException:
        
        # desperate situation
        numpy_image = GenerateNumPyImage( path, mime )
        
        if len( numpy_image.shape ) == 3:
            
            ( height, width, depth ) = numpy_image.shape
            
        else:
            
            ( height, width ) = numpy_image.shape
            
        
    
    width = max( width, 1 )
    height = max( height, 1 )
    
    return ( width, height )
    

def GetResolutionNumPy( numpy_image ):
    
    ( image_height, image_width, depth ) = numpy_image.shape
    
    return ( image_width, image_height )
    

THUMBNAIL_SCALE_DOWN_ONLY = 0
THUMBNAIL_SCALE_TO_FIT = 1
THUMBNAIL_SCALE_TO_FILL = 2

thumbnail_scale_str_lookup = {
    THUMBNAIL_SCALE_DOWN_ONLY : 'scale down only',
    THUMBNAIL_SCALE_TO_FIT : 'scale to fit',
    THUMBNAIL_SCALE_TO_FILL : 'scale to fill'
}

def GetThumbnailResolution( image_resolution: typing.Tuple[ int, int ], bounding_dimensions: typing.Tuple[ int, int ], thumbnail_scale_type: int, thumbnail_dpr_percent: int ) -> typing.Tuple[ int, int ]:
    
    ( im_width, im_height ) = image_resolution
    ( bounding_width, bounding_height ) = bounding_dimensions
    
    if thumbnail_dpr_percent != 100:
        
        thumbnail_dpr = thumbnail_dpr_percent / 100
        
        bounding_height = int( bounding_height * thumbnail_dpr )
        bounding_width = int( bounding_width * thumbnail_dpr )
        
    
    # this is appropriate for the crazy (0x0) svg or whatever we have, since we will _try_ to render it properly later on
    # but if it fails to render, we'll still get a fairly nice filetype.png or hydrus.png fallback
    # we don't want to pass around 0x0 and have a handler everywhere
    if im_width is None or im_width == 0 or im_height is None or im_height == 0:
        
        im_width = bounding_width
        im_height = bounding_width
        
    
    # TODO SVG thumbs should always scale up to the bounding dimensions
    
    if thumbnail_scale_type == THUMBNAIL_SCALE_DOWN_ONLY:
        
        if bounding_width >= im_width and bounding_height >= im_height:
            
            return ( im_width, im_height )
            
        
    
    image_ratio = im_width / im_height
    
    width_ratio = im_width / bounding_width
    height_ratio = im_height / bounding_height
    
    image_is_wider_than_bounding_box = width_ratio > height_ratio
    image_is_taller_than_bounding_box = height_ratio > width_ratio
    
    thumbnail_width = bounding_width
    thumbnail_height = bounding_height
    
    if thumbnail_scale_type in ( THUMBNAIL_SCALE_DOWN_ONLY, THUMBNAIL_SCALE_TO_FIT ):
        
        if image_is_taller_than_bounding_box: # i.e. the height will be at bounding height
            
            thumbnail_width = im_width / height_ratio
            
        elif image_is_wider_than_bounding_box: # i.e. the width will be at bounding width
            
            thumbnail_height = im_height / width_ratio
            
        
    elif thumbnail_scale_type == THUMBNAIL_SCALE_TO_FILL:
        
        # we do min 5.0 here to stop really tall and thin images getting zoomed in from width 1px to 150 and getting a thumbnail with a height of 75,000 pixels
        # in this case the line image is already crazy distorted, so we don't mind squishing it
        
        if image_is_taller_than_bounding_box: # i.e. the width will be at bounding width, the height will spill over
            
            thumbnail_height = bounding_width * min( 5.0, 1 / image_ratio )
            
        elif image_is_wider_than_bounding_box: # i.e. the height will be at bounding height, the width will spill over
            
            thumbnail_width = bounding_height * min( 5.0, image_ratio )
            
        
        # old stuff that actually clipped the size of the thing
        '''
        clip_x = 0
        clip_y = 0
        clip_width = im_width
        clip_height = im_height
        
        if width_ratio > height_ratio:
            
            clip_width = max( int( im_width * height_ratio / width_ratio ), 1 )
            clip_x = ( im_width - clip_width ) // 2
            
        elif height_ratio > width_ratio:
            
            clip_height = max( int( im_height * width_ratio / height_ratio ), 1 )
            clip_y = ( im_height - clip_height ) // 2
            
        
        clip_rect = ( clip_x, clip_y, clip_width, clip_height )
        '''
        
    
    thumbnail_width = int( thumbnail_width )
    thumbnail_height = int( thumbnail_height )
    
    thumbnail_width = max( thumbnail_width, 1 )
    thumbnail_height = max( thumbnail_height, 1 )
    
    return ( thumbnail_width, thumbnail_height )
    

def IsDecompressionBomb( path ) -> bool:
    
    # there are two errors here, the 'Warning' and the 'Error', which atm is just a test vs a test x 2 for number of pixels
    # 256MB bmp by default, ( 1024 ** 3 ) // 4 // 3
    # we'll set it at 512MB, and now catching error should be about 1GB
    
    PILImage.MAX_IMAGE_PIXELS = ( 512 * ( 1024 ** 2 ) ) // 3
    
    warnings.simplefilter( 'error', PILImage.DecompressionBombError )
    
    try:
        
        HydrusImageOpening.RawOpenPILImage( path )
        
    except ( PILImage.DecompressionBombError ):
        
        return True
        
    except:
        
        # pil was unable to load it, which does not mean it was a decomp bomb
        return False
        
    finally:
        
        PILImage.MAX_IMAGE_PIXELS = None
        
        warnings.simplefilter( 'ignore', PILImage.DecompressionBombError )
        
    
    return False
    

def ResizeNumPyImage( numpy_image: numpy.array, target_resolution, forced_interpolation = None ) -> numpy.array:
    
    ( target_width, target_height ) = target_resolution
    ( image_width, image_height ) = GetResolutionNumPy( numpy_image )
    
    if target_width == image_width and target_height == target_width:
        
        return numpy_image
        
    elif target_width > image_height or target_height > image_width:
        
        interpolation = cv2.INTER_LANCZOS4
        
    else:
        
        interpolation = cv2.INTER_AREA
        
    
    if forced_interpolation is not None:
        
        interpolation = forced_interpolation
        
    
    return cv2.resize( numpy_image, ( target_width, target_height ), interpolation = interpolation )
    
