import atexit
from _operator import itemgetter
from functools import partial, lru_cache
from itertools import chain, filterfalse
from operator import itemgetter
from typing import List, Union, Iterable

import fitz
import numpy as np
from PIL import Image
from funcy import merge, compose, rcompose, keep, lkeep
from pymonad.either import Right, Left, Either
from pymonad.tools import curry, identity

from image_prediction.exceptions import InvalidBox, BadXref
from image_prediction.formatter.formatters.enum import EnumFormatter
from image_prediction.image_extractor.extractor import ImageExtractor, ImageMetadataPair
from image_prediction.info import Info
from image_prediction.stitching.stitching import stitch_pairs
from image_prediction.stitching.utils import validate_box
from image_prediction.utils import get_logger
from image_prediction.utils.generic import bottom, left, right

logger = get_logger()


Doc = fitz.fitz.Document
Pag = fitz.fitz.Page
Img = Image.Image


class ParsablePDFImageExtractor(ImageExtractor):
    def __init__(self, verbose=False, tolerance=0):
        """

        Args:
            verbose: Whether to show progressbar
            tolerance: The tolerance in pixels for the distance between images, beyond which they will not be stitched
                together
        """
        self.doc: Union[Doc, None] = None
        self.verbose = verbose
        self.tolerance = tolerance

    def extract(self, pdf: bytes, page_range: range = None):
        self.doc = fitz.Document(stream=pdf)

        pages = extract_pages(self.doc, page_range) if page_range else self.doc

        image_metadata_pairs = chain.from_iterable(map(self.__process_images_on_page, pages))

        yield from image_metadata_pairs

    def __process_images_on_page(self, page: Pag):
        metadata = extract_valid_metadata(self.doc, page)

        either_image_metadata_pair_or_error_per_image = map(self.__metadatum_to_image_metadata_pair, metadata)
        valid_image_metadata_pairs = lkeep(take_good_log_bad, either_image_metadata_pair_or_error_per_image)
        valid_image_metadata_pairs_stitched = stitch_pairs(valid_image_metadata_pairs, tolerance=self.tolerance)

        clear_caches()

        yield from valid_image_metadata_pairs_stitched

    def __metadatum_to_image_metadata_pair(self, metadatum: dict) -> Either:
        return metadatum_to_image_metadata_pair(self.doc, metadatum)


def take_good_log_bad(pair: Either) -> Union[ImageMetadataPair, None]:
    return pair.either(log_error_context, identity)


def log_error_context(context: dict) -> None:
    logger.warning(f"Skipping bad image. {format_context(context)}")
    return None


def format_context(context: dict) -> str:
    return f"Reason: {context['reason'].rstrip('.')}. Metadata: {EnumFormatter()(context['metadata'])}"


def extract_pages(doc: Doc, page_range: range):
    page_range = range(page_range.start + 1, page_range.stop + 1)
    pages = map(doc.load_page, page_range)

    yield from pages


def validate_image(image: Img) -> Either:
    try:
        # TODO: stand-in heuristic for testing if image is valid => find cleaner solution (RED-5148)
        image.resize((100, 100)).convert("RGB")
        return Right(image)
    except (OSError, Exception):
        logger.warning(f"Invalid image encountered.")
        return Left("Invalid image.")


def extract_valid_metadata(doc: Doc, page: Pag) -> List[dict]:
    return compose(
        list,
        partial(add_alpha_channel_info, doc),
        filter_valid_metadata,
        get_metadata_for_images_on_page,
    )(page)


def metadatum_to_image_metadata_pair(doc: Doc, metadatum: dict) -> Either:
    image: Either = xref_to_image(doc, metadatum[Info.XREF]).bind(validate_image)
    image_metadata_pair: Either = make_eithered_image_metadata_pair(image, Right(metadatum))
    return image_metadata_pair


def add_alpha_channel_info(doc: Doc, metadata: Iterable[dict]) -> Iterable[dict]:
    def add_alpha_value_to_metadatum(metadatum: dict) -> dict:
        alpha = metadatum_to_alpha_value(metadatum)
        return {**metadatum, Info.ALPHA: alpha}

    xref_to_alpha = partial(has_alpha_channel, doc)
    metadatum_to_alpha_value = compose(xref_to_alpha, itemgetter(Info.XREF))

    yield from map(add_alpha_value_to_metadatum, metadata)


def filter_valid_metadata(metadata: Iterable[dict]) -> Iterable[dict]:
    yield from compose(filter_out_tiny_images, filter_out_invalid_metadata)(metadata)


def get_metadata_for_images_on_page(page: fitz.Page) -> Iterable[dict]:
    metadata = map(get_image_metadata, get_image_infos(page))
    metadata = add_page_metadata(page, metadata)

    yield from metadata


@lru_cache(maxsize=None)
def xref_to_image(doc: Doc, xref: int) -> Either:
    try:
        return Right(extract_image(doc, xref))
    except BadXref:
        return Left("Bad xref.")


def make_eithered_image_metadata_pair(image: Either, metadata: Either) -> Either:
    """Reference: haskell.org/tutorial/monads.html"""

    def context(value):
        return {"reason": value, "metadata": metadata.either(bottom, identity)}

    # What is most readable?

    # 1)
    # return Either.apply(make_image_metadata_pair).to_arguments(image, metadata).either(left(context), right(identity))

    # 2)
    # m (a -> b -> c) -> m a -> m b -> m c
    # return Right(make_image_metadata_pair).amap(image).amap(metadata).either(left(context), right(identity))

    # 3)
    # a := Image
    # b := Metadata
    # c := ImageMetadataPair
    return (
        image.bind(right(make_image_metadata_pair))  # m a >>= m (a -> b -> c) -> m (b -> c)
        .amap(metadata)  # m (b -> c) <*> m b -> m c
        .either(
            left(context),
            right(identity),
        )
    )


@curry(2)
def make_image_metadata_pair(image: Img, metadatum: dict) -> ImageMetadataPair:
    return ImageMetadataPair(image, metadatum)


def extract_image(doc: Doc, xref: int) -> Img:
    try:
        pixmap = fitz.Pixmap(doc, xref)
    except ValueError as err:
        msg = f"Xref {xref} is invalid, skipping extraction."
        logger.debug(msg)
        raise BadXref(msg) from err

    array = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape((pixmap.h, pixmap.w, pixmap.n))
    array = normalize_channels(array)
    return Image.fromarray(array)


def has_alpha_channel(doc: Doc, xref: int):

    maybe_image = load_image_handle_from_xref(doc, xref)
    maybe_smask = maybe_image["smask"] if maybe_image else None

    if maybe_smask:  # Use monad
        return any([doc.extract_image(maybe_smask) is not None, bool(fitz.Pixmap(doc, maybe_smask).alpha)])
    else:
        try:
            return bool(fitz.Pixmap(doc, xref).alpha)
        except ValueError:
            logger.debug(f"Encountered invalid xref `{xref}` in {doc.metadata.get('title', '<no title>')}.")
            return False


def filter_out_tiny_images(metadata: Iterable[dict]) -> Iterable[dict]:
    yield from filterfalse(tiny, metadata)


def filter_out_invalid_metadata(metadata: Iterable[dict]) -> Iterable[dict]:
    def __validate_box(box):
        try:
            return validate_box(box)
        except InvalidBox as err:
            logger.debug(f"Dropping invalid metadatum, reason: {err}")

    yield from keep(__validate_box, metadata)


def get_image_metadata(image_info: dict) -> dict:

    xref, coords = itemgetter("xref", "bbox")(image_info)
    x1, y1, x2, y2 = map(rounder, coords)

    width = abs(x2 - x1)
    height = abs(y2 - y1)

    return {
        Info.WIDTH: width,
        Info.HEIGHT: height,
        Info.X1: x1,
        Info.X2: x2,
        Info.Y1: y1,
        Info.Y2: y2,
        Info.XREF: xref,
    }


@lru_cache(maxsize=None)
def get_image_infos(page: Pag) -> List[dict]:
    return page.get_image_info(xrefs=True)


def add_page_metadata(page: Pag, metadata: Iterable[dict]) -> Iterable[dict]:
    yield from map(partial(merge, get_page_metadata(page)), metadata)


def normalize_channels(array: np.ndarray):
    if not array.ndim == 3:
        array = np.expand_dims(array, axis=-1)

    if array.shape[-1] == 4:
        array = array[..., :3]
    elif array.shape[-1] == 1:
        array = np.concatenate([array, array, array], axis=-1)
    elif array.shape[-1] != 3:
        logger.warning(f"Unexpected image format: {array.shape}.")
        raise ValueError(f"Unexpected image format: {array.shape}.")

    return array


@lru_cache(maxsize=None)
def load_image_handle_from_xref(doc: Doc, xref: int) -> Union[dict, None]:  # TODO: use Monad
    return doc.extract_image(xref)


def get_page_metadata(page: Pag) -> dict:
    page_width, page_height = map(rounder, page.mediabox_size)

    return {
        Info.PAGE_WIDTH: page_width,
        Info.PAGE_HEIGHT: page_height,
        Info.PAGE_IDX: page.number,
    }


rounder = rcompose(round, int)


def tiny(metadatum: dict) -> bool:
    return metadatum[Info.WIDTH] * metadatum[Info.HEIGHT] <= 4


def clear_caches() -> None:
    get_image_infos.cache_clear()
    load_image_handle_from_xref.cache_clear()
    xref_to_image.cache_clear()


atexit.register(clear_caches)