from __future__ import annotations
from dataclasses import dataclass
import hashlib
import os
from os.path import join
import math
import shutil
from typing import Optional
import sewar
import numpy as np
from PIL import Image
import png

from subprocess import CompletedProcess


class ThumbInvalid(Exception):
    def __init__(self, original_path, thumb_path, msg):
        Exception.__init__(self)
        self.original_path = original_path
        self.thumb_path = thumb_path
        self.msg = msg

    def __str__(self):
        return "Invalid thumb: {m}\nOriginal:{o}\nThumb:{t}".format(
            m=self.msg,
            o=self.original_path,
            t=self.thumb_path,
        )


class ThumbValidator:
    """
    Notes:
    
    * https://towardsdatascience.com/measuring-similarity-in-two-images-using-python-b72233eb53c6
    * https://pypng.readthedocs.io/en/latest/index.html
    """
    @staticmethod
    def validate_case(case: Case):
        for size, length, actual in zip(case.sizes, case.lengths, case.thumb_paths):
            expected = case.test_images["thumb_" + size]
            validator = ThumbValidator(expected, actual)
            validator.validate(length)

    @dataclass
    class PNGData:
        width: int
        height: int
        info: dict
        reader: png.Reader

        @staticmethod
        def from_path(path) -> ThumbValidator.PNGData:
            r = png.Reader(path)
            t = r.read()
            return ThumbValidator.PNGData(
                width=t[0],
                height=t[1],
                info=t[3],
                reader=r
            )

    def __init__(self, path_of_expected_image: str, path_of_actual_image: str):
        self.path_of_expected_image = path_of_expected_image
        self.path_of_actual_image = path_of_actual_image

    def _raise(self, msg: str):
        raise ThumbInvalid(
            original_path=self.path_of_expected_image,
            thumb_path=self.path_of_actual_image,
            msg=msg
        )

    def _validate_png_data(self, png_data: ThumbValidator.PNGData, expected_max_length: int):
        # might actually be a `> length`
        # see https://gitlab.com/DLF/allmytoes/-/issues/15
        if max(png_data.width, png_data.height) != expected_max_length:
            self._raise(
                "Thumb size is {}x{} but should be max. {}".format(
                    png_data.width,
                    png_data.height,
                    expected_max_length,
                )
            )

        if png_data.info['bitdepth'] != 8:
            self._raise("Thumb must be 8 bit but is {}.".format(
                png_data.info['bitdepth']
            ))
        if png_data.info['interlace'] != 0:
            self._raise("Thumb must be non-interlaced, but interlace is {}.".format(
                png_data.info['interlace']
            ))

    def _validate_image_similarity(self, expected_max_length: int):
        actual_pil_image = Image.open(self.path_of_actual_image)
        actual_array = np.asarray(actual_pil_image)
        expected_pil_image = Image.open(self.path_of_expected_image)
        maxsize = (expected_max_length, expected_max_length)
        #expected_pil_image.thumbnail(maxsize, Image.ANTIALIAS)
        expected_array = np.asarray(expected_pil_image)
        deviation = sewar.mse(expected_array, actual_array)
        max_deviation = math.ceil(math.sqrt(expected_max_length)) * 1.1
        if not deviation < max_deviation:
            self._raise(
                "Thumb does not look equal. Expected like: {}, Generated thumb: {}, deviation: {}, allowed deviation: {}".format(
                self.path_of_expected_image,
                self.path_of_actual_image,
                deviation,
                max_deviation,
            )
        )


    def validate(self, expected_max_side_length: int):
        if not os.path.isfile(self.path_of_actual_image):
            self._raise("Actual image does not exist.")
        if not os.path.isfile(self.path_of_expected_image):
            self._raise("Expected image example does not exist.")
        png_data = ThumbValidator.PNGData.from_path(self.path_of_actual_image)
        self._validate_png_data(png_data, expected_max_side_length)
        self._validate_image_similarity(expected_max_side_length)



class Case:
    def __init__(self, working_dir: str, test_images: dict[str, str]):
        # The temporary working dir for the scenario, also acting as user-home dir.
        self.working_dir = working_dir

        # A dict of prepared images for testing. Contains input files and "expected" images for comparison.
        self.test_images = test_images

        # The infile used for a test case.
        self.infile: Optional[str] = None

        # The mtime of the infile when set.
        self.infile_mtime: Optional[int] = None

        # The cache dir
        self.cache_dir = join(self.working_dir, ".cache")

        # Default thumb directory.
        self.thumb_dir = join(self.cache_dir, "thumbnails")

        # All expected sizes, named like the size directories.
        self.sizes = ["normal", "large", "x-large", "xx-large"]

        # All max edge lengths for the respective sizes
        self.lengths = [128, 256, 512, 1024]

        # Map paths of thumbs created by the test suite to their mtime after creation.
        # This allows to test later if the thumbs have been touched.
        self.prepared_thumbs_mtimes = {}

        # The result of running AllMyToes
        self.result: Optional[CompletedProcess] = None

    def copy_and_set_infile(self, source_path: str):
        infile_name = "infile." + source_path.split(".")[-1]
        self.infile = join(self.working_dir, infile_name)
        shutil.copyfile(source_path, self.infile)
        self.infile_mtime = int(os.path.getmtime(self.infile))

    def hard_set_infile(self, infile_path: str):
        self.infile = infile_path

    @property
    def infile_hash(self) -> str:
        """
        The hash for the input file like used for the thumb-naming.
        """
        assert self.infile
        infile_uri = "file://{}".format(self.infile)
        return hashlib.md5(infile_uri.encode("utf-8")).hexdigest()

    @property
    def thumb_name(self) -> str:
        return self.infile_hash + ".png"

    @property
    def thumb_paths(self) -> list[str]:
        """
        List of the thumb dirs for the different sizes.

        Creates those in the test-directory on the fly if not existing.
        """
        assert self.infile
        thumb_name = self.thumb_name
        result = []
        for size in self.sizes:
            d = join(self.thumb_dir, size)
            if not os.path.exists(d):
                os.makedirs(d)
            path = join(d, thumb_name)
            result.append(path)
        return result


