from collections.abc import Iterable
from dataclasses import dataclass
from enum import Enum, auto
from pathlib import Path
from typing import TypeAlias

import PIL.Image
import pytest

FilePathOrPILImage: TypeAlias = str | Path | PIL.Image.Image

from hordelib.utils.distance import (
    CosineSimilarityResult,
    CosineSimilarityResultCode,
    HistogramDistanceResult,
    HistogramDistanceResultCode,
    evaluate_image_distance,
    is_cosine_similarity_fail,
    is_histogram_distance_fail,
)


class ImageSimilarityDefaultsTypes(Enum):
    STANDARD_INFERENCE = auto()
    LORA = auto()


@dataclass
class ImageSimilarityConstraints:
    cosine_fail_floor: CosineSimilarityResultCode
    cosine_warn_floor: CosineSimilarityResultCode
    histogram_fail_threshold: HistogramDistanceResultCode
    histogram_warn_threshold: HistogramDistanceResultCode


INFERENCE_SIMILARITY_DEFAULTS = ImageSimilarityConstraints(
    cosine_fail_floor=CosineSimilarityResultCode.PARTIALLY_SIMILAR,
    cosine_warn_floor=CosineSimilarityResultCode.CONSIDERABLY_SIMILAR,
    histogram_fail_threshold=HistogramDistanceResultCode.SKIP,
    histogram_warn_threshold=HistogramDistanceResultCode.SKIP,
)

LORA_SIMILARITY_DEFAULTS = ImageSimilarityConstraints(
    cosine_fail_floor=CosineSimilarityResultCode.SKIP,
    cosine_warn_floor=CosineSimilarityResultCode.CONSIDERABLY_SIMILAR,
    histogram_fail_threshold=HistogramDistanceResultCode.SKIP,
    histogram_warn_threshold=HistogramDistanceResultCode.SKIP,
)


class ImageSimilarityResultCode(Enum):
    FAIL = 0
    SKIP = 1
    PASS = 2


@dataclass
class ImageSimilarityResult:
    result_code: ImageSimilarityResultCode
    cosine_similarity_result: CosineSimilarityResult | float
    histogram_distance_result: HistogramDistanceResult | float


def check_image_similarity_pytest(
    img1: FilePathOrPILImage,
    img2: FilePathOrPILImage,
    *,
    similarity_constraints: ImageSimilarityConstraints,
) -> ImageSimilarityResult:
    cosine_similarity_result, histogram_distance_result = evaluate_image_distance(img1, img2)

    cosine_fail = (
        is_cosine_similarity_fail(
            result_to_check=cosine_similarity_result,
            threshold=similarity_constraints.cosine_fail_floor,
        )
        if similarity_constraints.cosine_fail_floor != CosineSimilarityResultCode.SKIP
        else False
    )
    cosine_warn = (
        is_cosine_similarity_fail(
            result_to_check=cosine_similarity_result,
            threshold=similarity_constraints.cosine_warn_floor,
        )
        if similarity_constraints.cosine_warn_floor != CosineSimilarityResultCode.SKIP
        else False
    )

    histogram_fail = (
        is_histogram_distance_fail(
            result_to_check=histogram_distance_result,
            maximum=similarity_constraints.histogram_fail_threshold,
        )
        if similarity_constraints.histogram_fail_threshold != HistogramDistanceResultCode.SKIP
        else False
    )
    histogram_warn = (
        is_histogram_distance_fail(
            result_to_check=histogram_distance_result,
            maximum=similarity_constraints.histogram_warn_threshold,
        )
        if similarity_constraints.histogram_warn_threshold != HistogramDistanceResultCode.SKIP
        else False
    )

    if cosine_fail or histogram_fail:
        return ImageSimilarityResult(
            ImageSimilarityResultCode.FAIL,
            cosine_similarity_result,
            histogram_distance_result,
        )
    if cosine_warn or histogram_warn:
        return ImageSimilarityResult(
            ImageSimilarityResultCode.SKIP,
            cosine_similarity_result,
            histogram_distance_result,
        )

    return ImageSimilarityResult(
        ImageSimilarityResultCode.PASS,
        cosine_similarity_result,
        histogram_distance_result,
    )


def do_pytest_action(result: ImageSimilarityResult) -> bool:
    """Perform the appropriate pytest action (e.g. `pytest.fail(...)`) based on the image similarity result.

    Args:
        result (ImageSimilarityResult): _description_

    Returns:
        bool: _description_
    """
    if ImageSimilarityResult.result_code == ImageSimilarityResultCode.FAIL:
        pytest.fail(
            f"Image similarity check failed:\n"
            f"cosine_similarity_result={ImageSimilarityResult.cosine_similarity_result},\n"
            f"histogram_distance_result={ImageSimilarityResult.histogram_distance_result}",
        )
    elif ImageSimilarityResult.result_code == ImageSimilarityResultCode.SKIP:
        pytest.skip(
            f"Image similarity check skipped:\n"
            f"cosine_similarity_result={ImageSimilarityResult.cosine_similarity_result},\n"
            f"histogram_distance_result={ImageSimilarityResult.histogram_distance_result}",
        )

    return True


def check_single_inference_image_similarity(
    img1: FilePathOrPILImage,
    img2: FilePathOrPILImage,
) -> bool:
    return check_list_inference_images_similarity(
        [(img1, img2)],
    )


def check_single_lora_image_similarity(
    img1: FilePathOrPILImage,
    img2: FilePathOrPILImage,
) -> bool:
    return check_list_lora_images_similarity(
        [(img1, img2)],
    )


def check_list_lora_images_similarity(
    list_of_image_pairs: Iterable[tuple[FilePathOrPILImage, FilePathOrPILImage]],
) -> bool:
    return check_list_images_similarity(
        list_of_image_pairs,
        similarity_constraints=LORA_SIMILARITY_DEFAULTS,
    )


def check_list_inference_images_similarity(
    list_of_image_pairs: Iterable[tuple[FilePathOrPILImage, FilePathOrPILImage]],
) -> bool:
    return check_list_images_similarity(
        list_of_image_pairs,
        similarity_constraints=INFERENCE_SIMILARITY_DEFAULTS,
    )


def check_list_images_similarity(
    list_of_image_pairs: Iterable[tuple[FilePathOrPILImage, FilePathOrPILImage]],
    *,
    similarity_constraints: ImageSimilarityConstraints,
) -> bool:
    all_results: list[tuple[tuple[FilePathOrPILImage, FilePathOrPILImage], ImageSimilarityResult]] = []
    for img1, img2 in list_of_image_pairs:
        image_similarity_result = check_image_similarity_pytest(
            img1,
            img2,
            similarity_constraints=similarity_constraints,
        )
        all_results.append(((img1, img2), image_similarity_result))

    if all(result[1].result_code == ImageSimilarityResultCode.PASS for result in all_results):
        return True

    all_failed_results = [result for result in all_results if result[1].result_code == ImageSimilarityResultCode.FAIL]
    all_skipped_results = [result for result in all_results if result[1].result_code == ImageSimilarityResultCode.SKIP]

    complete_error_message = ""

    if len(all_failed_results) > 0:
        for img_pair, result in all_failed_results:
            complete_error_message += (
                f"Image similarity check failed:\n"
                f"cosine_similarity_result={result.cosine_similarity_result},\n"
                f"histogram_distance_result={result.histogram_distance_result}\n"
            )
            complete_error_message += f"img1={img_pair[0]}\n"
            complete_error_message += f"img2={img_pair[1]}\n"
        complete_error_message += "all_failed_results:\n"
    if len(all_skipped_results) > 0:
        for img_pair, result in all_skipped_results:
            complete_error_message += (
                f"Image similarity check skipped:\n"
                f"cosine_similarity_result={result.cosine_similarity_result},\n"
                f"histogram_distance_result={result.histogram_distance_result}\n"
            )
            complete_error_message += f"img1={img_pair[0]}\n"
            complete_error_message += f"img2={img_pair[1]}\n"
    if len(all_failed_results) > 0:
        pytest.fail(complete_error_message)
    elif len(all_skipped_results) > 0:
        pytest.skip(complete_error_message)

    return True
