from itertools import chain
from operator import itemgetter

import fitz

from image_prediction.image_extractor.extractor import ImageExtractor, ImageMetadataPair


class ParsablePDFImageExtractor(ImageExtractor):
    def __init__(self):
        self.doc: fitz.fitz.Document = None
        self.page: fitz.fitz.Page = None

    def __build_metadata(self, xref):
        metadata = self.page.get_image_info(xref)
        page_width, page_height = self.page.mediabox_size
        metadata = {**metadata, "page_width": page_width, "page_height": page_height, "page_idx": self.page.number}
        return metadata

    def __process_images_on_page(self, page: fitz.fitz.Page):
        self.page = page

        image_handles = page.get_images(full=True)
        xrefs = itemgetter(0)(image_handles)
        images = map(lambda xref: self.doc.extract_image(xref)["image"], xrefs)
        metadata = map(self.__build_metadata, xrefs)
        return map(ImageMetadataPair, zip(images, metadata))

    def extract(self, pdf: bytes):
        self.doc = fitz.Document(stream=pdf)
        image_metadata_pairs = chain(*map(self.__process_images_on_page, self.doc))
        return image_metadata_pairs