import json
from dataclasses import dataclass
from functools import partial
from operator import itemgetter
from typing import Literal, SupportsIndex, Tuple

import fitz  # type: ignore
import numpy as np
from funcy import compose, lfilter  # type: ignore
from kn_utils.logging import logger  # type: ignore
from numpy import ndarray as Array

BBoxType = tuple[int | float, int | float, int | float, int | float]


@dataclass
class PageInfo:
    page_num: int
    rotation_matrix: fitz.Matrix
    transformation_matrix: fitz.Matrix
    dpi: int
    width: int | float
    height: int | float
    image_width: int | float
    image_height: int | float
    rotation: int


def transform_image_coordinates_to_pdf_coordinates(
    bbox: BBoxType,
    rotation_matrix: fitz.Matrix,
    transformation_matrix: fitz.Matrix,
    dpi: int | None = None,
) -> Tuple:
    x1, y1, x2, y2 = map(lambda x: (x / dpi) * 72, bbox) if dpi else bbox  # Convert to points, can be done before
    rect = fitz.Rect(x1, y1, x2, y2)
    rect = rect * rotation_matrix * transformation_matrix

    return rect.x0, rect.y0, rect.x1, rect.y1


def rescale_to_pdf(bbox: BBoxType, page_info: PageInfo) -> BBoxType:
    round3 = lambda x: tuple(map(lambda y: round(y, 3), x))

    pdf_h, pdf_w = page_info.height, page_info.width

    if page_info.rotation in {90, 270}:
        pdf_h, pdf_w = pdf_w, pdf_h

    pix_h, pix_w = page_info.image_height, page_info.image_width
    ratio_h, ratio_w = pdf_h / pix_h, pdf_w / pix_w
    ratio_w, ratio_h, pdf_w, pdf_h, pix_w, pix_h = round3((ratio_w, ratio_h, pdf_w, pdf_h, pix_w, pix_h))

    return round3((bbox[0] * ratio_w, bbox[1] * ratio_h, bbox[2] * ratio_w, bbox[3] * ratio_h))


def mirror_horizontal(bbox: BBoxType, page_height: int | float) -> BBoxType:
    x0, y0, x1, y1 = bbox
    y0_new = page_height - y0
    y1_new = page_height - y1

    return x0, y0_new, x1, y1_new


def mirror_vertical(bbox: BBoxType, page_width: int | float) -> BBoxType:
    x0, y0, x1, y1 = bbox
    x0_new = page_width - x1
    x1_new = page_width - x0

    return x0_new, y0, x1_new, y1


def derotate_image(bbox: BBoxType, page_info: PageInfo) -> BBoxType:

    logger.debug(f"{page_info.rotation=}")
    match page_info.rotation:
        case 0:
            bbox = mirror_horizontal(bbox, page_info.height)
        case 90:
            pass
        case 180:
            bbox = mirror_vertical(bbox, page_info.height)
        case 270:
            bbox = mirror_vertical(mirror_horizontal(bbox, page_info.height), page_info.height)
        case _:
            logger.warning(f"Unknown rotation: {page_info.rotation}")
            pass

    return bbox


def transform_table_lines_by_page_info(bboxes: dict, offsets: tuple, page_info: PageInfo) -> dict:

    transform = partial(rescale_to_pdf, page_info=page_info)
    logger.debug(f"{offsets=}")

    def apply_offsets(line: tuple) -> tuple:
        x1, y1, x2, y2 = line
        offset_x, offset_y = offsets
        offset_y = page_info.height - offset_y
        logger.debug((f"new offsets: {offset_x}, {offset_y}"))

        return (x1 + offset_x, y1 + offset_y, x2 + offset_x, y2 + offset_y)

    derotate = partial(derotate_image, page_info=page_info)

    unpack = itemgetter("x1", "y1", "x2", "y2")
    pack = lambda x: {"x1": x[0], "y1": x[1], "x2": x[2], "y2": x[3]}
    convert = compose(pack, derotate, apply_offsets, transform, unpack)

    table_lines = bboxes.get("tableLines", [])
    bboxes["tableLines"] = list(map(convert, table_lines))

    for i in range(len(table_lines)):
        logger.debug(json.dumps(table_lines[i], indent=4))
        logger.debug(json.dumps(bboxes["tableLines"][i], indent=4))

    return bboxes


def extract_images_from_pdf(
    pdf_bytes: bytes, vlp_output: dict, dpi: int = 200
) -> tuple[list[Array], list[dict], list[PageInfo]]:
    with fitz.open(stream=pdf_bytes) as fh:
        table_images = []
        table_info = []
        page_info = []

        vlp_output = vlp_output["data"] if isinstance(vlp_output, dict) else vlp_output

        for page_dict in vlp_output:
            page_num = int(page_dict["page_idx"])
            boxes = page_dict["boxes"]
            boxes = filter(lambda box_obj: box_obj["label"] == "table", boxes)

            page: fitz.Page = fh[page_num]
            page.wrap_contents()

            page_image = page.get_pixmap(dpi=200)
            current_page_info = PageInfo(
                page_num,
                page.rotation_matrix,
                page.transformation_matrix,
                dpi,
                page.rect[-2],
                page.rect[-1],
                page_image.w,
                page_image.h,
                page.rotation,
            )

            for box_obj in boxes:
                bbox = box_obj["box"]
                x1, y1, x2, y2 = itemgetter("x1", "y1", "x2", "y2")(bbox)
                rect = fitz.Rect((x1, y1), (x2, y2))
                # FIXME: Check if de-rotation works as intended and is necessary at all.
                #  Note that there exists also a derotation_matrix. If changing this, also change the
                #  current_page_info object to include the derotation_matrix.
                rect = rect * page.transformation_matrix * page.rotation_matrix
                pixmap = page.get_pixmap(clip=rect, dpi=dpi, colorspace=fitz.csGRAY)
                shape = (pixmap.h, pixmap.w, pixmap.n) if pixmap.n > 1 else (pixmap.h, pixmap.w)
                image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(*shape)

                table_images.append(image)
                table_info.append(
                    {
                        "pageNum": page_num,
                        "bbox": bbox,
                        "uuid": box_obj["uuid"],
                        "label": box_obj["label"],
                    }
                )
                page_info.append(current_page_info)

    return table_images, table_info, page_info