import base64

import numpy as np
import pdfminer.pdfinterp
import pymupdf
from pdfminer.layout import LTChar, LTFigure
from pdfminer.pdffont import PDFCIDFont, PDFFont

from yadt.doclayout import DocLayoutModel
from yadt.document_il import il_version_1
from yadt.translation_config import TranslationConfig


class ILCreater:
    def __init__(self, translation_config: TranslationConfig):
        self.current_page: il_version_1.Page = None
        self.mupdf: pymupdf.Document = None
        self.model = DocLayoutModel.load_available()
        self.docs = il_version_1.Document(page=[])
        self.stroking_color_space_name = None
        self.non_stroking_color_space_name = None
        self.translation_config = translation_config

    # pdf32000 page 171
    def on_stroking_color_space(self, color_space_name):
        self.stroking_color_space_name = color_space_name

    def on_non_stroking_color_space(self, color_space_name):
        self.non_stroking_color_space_name = color_space_name

    def on_page_start(self):
        self.current_page = il_version_1.Page(
            pdf_font=[],
            pdf_character=[],
            page_layout=[],
            # currently don't support UserUnit page parameter
            # pdf32000 page 79
            unit="point",
        )
        self.current_page_font_name_id_map = {}
        self.docs.page.append(self.current_page)

    def on_page_crop_box(
        self,
        x0: float | int,
        y0: float | int,
        x1: float | int,
        y1: float | int,
    ):
        box = il_version_1.Box(
            x=float(x0), y=float(y0), x2=float(x1), y2=float(y1)
        )
        self.current_page.cropbox = il_version_1.Cropbox(box=box)

    def on_page_media_box(
        self,
        x0: float | int,
        y0: float | int,
        x1: float | int,
        y1: float | int,
    ):
        box = il_version_1.Box(
            x=float(x0), y=float(y0), x2=float(x1), y2=float(y1)
        )
        self.current_page.mediabox = il_version_1.Mediabox(box=box)

    def on_page_number(self, page_number: int):
        assert isinstance(page_number, int)
        assert page_number >= 0
        self.current_page.page_number = page_number

        self.on_page_layout(page_number)

    def on_page_base_operation(self, operation: str):
        self.current_page.base_operations = il_version_1.BaseOperations(
            value=operation
        )

    def on_page_resource_font(self, font: PDFFont, xref_id: int, font_id: str):
        font_name = font.fontname
        if isinstance(font.fontname, bytes):
            try:
                font_name = font.fontname.decode("utf-8")
            except UnicodeDecodeError:
                font_name = "BASE64:" + base64.b64encode(font_name).decode(
                    "utf-8"
                )
        encoding_length = 1
        if isinstance(font, PDFCIDFont):
            encoding_length = 2
        il_font_metadata = il_version_1.PdfFont(
            name=font_name,
            xref_id=xref_id,
            font_id=font_id,
            encoding_length=encoding_length,
        )
        self.current_page_font_name_id_map[font_name] = font_id
        self.current_page.pdf_font.append(il_font_metadata)

    def create_graphic_state(self, gs: pdfminer.pdfinterp.PDFGraphicState):
        graphic_state = il_version_1.GraphicState()
        for k, v in gs.__dict__.items():
            if v is None:
                continue
            if k in ["scolor", "ncolor"]:
                if isinstance(v, tuple):
                    v = list(v)
                else:
                    v = [v]
                setattr(graphic_state, k, v)
                continue
            if k == "linewidth":
                graphic_state.linewidth = float(v)
                continue
            continue
            raise NotImplementedError

        graphic_state.stroking_color_space_name = (
            self.stroking_color_space_name
        )
        graphic_state.non_stroking_color_space_name = (
            self.non_stroking_color_space_name
        )
        return graphic_state

    def on_lt_char(self, char: LTChar):
        gs = self.create_graphic_state(char.graphicstate)
        bbox = il_version_1.Box(
            char.bbox[0], char.bbox[1], char.bbox[2], char.bbox[3]
        )

        font_id = self.current_page_font_name_id_map[char.font.fontname]
        char_id = char.cid
        char_unicode = char.get_text()
        advance = char.adv
        if char.matrix[0] == 0 and char.matrix[3] == 0:
            vertical = True
        else:
            vertical = False
        pdf_style = il_version_1.PdfStyle(
            font_id=font_id,
            font_size=char.size,
            graphic_state=gs,
        )
        pdf_char = il_version_1.PdfCharacter(
            box=bbox,
            pdf_character_id=char_id,
            advance=advance,
            char_unicode=char_unicode,
            vertical=vertical,
            pdf_style=pdf_style,
        )
        self.current_page.pdf_character.append(pdf_char)

    def on_page_layout(self, page_number):
        if (
            self.translation_config.should_translate_page(page_number + 1)
            is False
        ):
            return
        pix = self.mupdf[page_number].get_pixmap()
        image = np.fromstring(pix.samples, np.uint8).reshape(
            pix.height, pix.width, 3
        )[:, :, ::-1]
        h, w = pix.height, pix.width
        layouts = self.model.predict(image, imgsz=int(pix.height / 32) * 32)[0]
        id = 0
        for layout in layouts.boxes:
            id += 1
            # Convert the coordinate system from the picture coordinate system to the il coordinate system
            x0, y0, x1, y1 = layout.xyxy
            x0, y0, x1, y1 = (
                np.clip(int(x0 - 1), 0, w - 1),
                np.clip(int(h - y1 - 1), 0, h - 1),
                np.clip(int(x1 + 1), 0, w - 1),
                np.clip(int(h - y0 + 1), 0, h - 1),
            )
            page_layout = il_version_1.PageLayout(
                id=id,
                box=il_version_1.Box(
                    x0.item(), y0.item(), x1.item(), y1.item()
                ),
                conf=layout.conf.item(),
                class_name=layouts.names[layout.cls],
            )
            self.docs.page[page_number].page_layout.append(page_layout)

    def create_il(self):
        pages = [
            page
            for page in self.docs.page
            if self.translation_config.should_translate_page(
                page.page_number + 1
            )
        ]
        self.docs.page = pages
        return self.docs

    def on_total_pages(self, total_pages: int):
        assert isinstance(total_pages, int)
        assert total_pages > 0
        self.docs.total_pages = total_pages

    def on_pdf_figure(self, figure: LTFigure):
        box = il_version_1.Box(
            figure.bbox[0], figure.bbox[1], figure.bbox[2], figure.bbox[3]
        )
        self.current_page.pdf_figure.append(il_version_1.PdfFigure(box=box))
