from dataclasses import asdict
from operator import itemgetter, truth
from typing import Callable, Generator

from funcy import flatten, lmap  # type: ignore
from pdf2img.conversion import convert_pages_to_images  # type: ignore
from pdf2img.default_objects.image import ImageInfo, ImagePlus  # type: ignore
from pdf2img.default_objects.rectangle import RectanglePlus  # type: ignore

from cv_analysis.figure_detection.figure_detection import detect_figures
from cv_analysis.table_inference import infer_lines
from cv_analysis.table_parsing import parse_lines, parse_tables
from cv_analysis.utils.image_extraction import (
    extract_images_from_pdf,
    transform_table_lines_by_page_info,
)
from cv_analysis.utils.structures import Rectangle


def get_analysis_pipeline(operation, table_parsing_skip_pages_without_images=True):
    if operation == "table":
        return make_analysis_pipeline(
            parse_lines,
            table_parsing_formatter,
            dpi=200,
            skip_pages_without_images=table_parsing_skip_pages_without_images,
        )
    if operation == "table_cells":
        return make_analysis_pipeline(
            parse_tables,
            table_parsing_cells_formatter,
            dpi=200,
            skip_pages_without_images=table_parsing_skip_pages_without_images,
        )
    if operation == "figure":
        return make_analysis_pipeline(detect_figures, figure_detection_formatter, dpi=200)
    if operation == "table_image_inference":  # TODO: fix pyinfra input
        return make_image_analysis_pipeline(
            infer_lines,
        )
    # else:
    # raise


def make_image_analysis_pipeline(
    analysis_fn,
) -> Callable[[dict], Generator[dict, bytes, None]]:
    def analyse_pipeline(data: dict) -> Generator[dict, bytes, None]:
        pdf_bytes = data["pdf"]
        vlp_output = data["vlp_output"]
        images, info, page_info = extract_images_from_pdf(pdf_bytes, vlp_output)
        # rel_bboxes = map()
        img_results = lmap(analysis_fn, images)

        def make_offsets(): ...

        offsets = map(itemgetter("x1", "y2"), map(itemgetter("bbox"), info))
        # print("before", img_results)
        img_results = lmap(transform_table_lines_by_page_info, img_results, offsets, page_info)
        # print("after", img_results)
        results = map(lambda i: info[i] | img_results[i], range(len(info)))

        yield from results

    return analyse_pipeline


def make_analysis_pipeline(analysis_fn, formatter, dpi, skip_pages_without_images=False):
    def analyse_pipeline(pdf: bytes, index=None):
        def parse_page(page: ImagePlus):
            image = page.asarray()
            rects = analysis_fn(image)
            if not rects:
                return None
            infos = formatter(rects, page, dpi)
            return infos

        pages = convert_pages_to_images(
            pdf,
            index=index,
            dpi=dpi,
            skip_pages_without_images=skip_pages_without_images,
        )
        results = map(parse_page, pages)

        yield from flatten(filter(truth, results))

    return analyse_pipeline


def table_parsing_formatter(lines: list[dict[str, float]], page: ImagePlus, dpi):
    return {"pageInfo": page.asdict(natural_index=True), "tableLines": lines}


def table_parsing_cells_formatter(rects, page: ImagePlus, dpi):
    def format_rect(rect: Rectangle):
        rect_plus = RectanglePlus.from_pixels(*rect.xyxy(), page.info, alpha=False, dpi=dpi)
        return rect_plus.asdict(derotate=True)

    bboxes = lmap(format_rect, rects)

    return {"pageInfo": page.asdict(natural_index=True), "tableCells": bboxes}


def figure_detection_formatter(rects, page, dpi):
    def format_rect(rect: Rectangle):
        rect_plus = RectanglePlus.from_pixels(*rect.xyxy(), page.info, alpha=False, dpi=dpi)
        return asdict(ImageInfo(page.info, rect_plus.asbbox(derotate=False), rect_plus.alpha))

    return lmap(format_rect, rects)
