"""
Copyright and License
----------------------
Copyright 2024 Artifex Software, Inc.
License GNU Affero GPL 3.0
"""
from pathlib import Path
from typing import List

import fitz

if fitz.pymupdf_version_tuple < (1, 24, 0):
    raise NotImplementedError("PyMuPDF version 1.24.0 or later is needed.")

def to_markdown(doc: fitz.Document, pages: list = None) -> str:
    """Process the document and return the text of its selected pages."""
    if isinstance(doc, str):
        doc = fitz.open(doc)
    if not pages:  # use all pages if argument not given
        pages = range(doc.page_count)

    md_string = ""

    for pno in pages:
        page = doc[pno]
        # 1. first locate all tables on page
        tabs = page.find_tables()

        # 2. make a list of table boundary boxes, sort by top-left corner.
        # Must include the header bbox, which may be external.
        tab_rects = sorted(
            [
                (fitz.Rect(t.bbox) | fitz.Rect(t.header.bbox), i)
                for i, t in enumerate(tabs.tables)
            ],
            key=lambda r: (r[0].y0, r[0].x0),
        )

        # 3. final list of all text and table rectangles
        text_rects = []
        # compute rectangles outside tables and fill final rect list
        for i, (r, idx) in enumerate(tab_rects):
            if i == 0:  # compute rect above all tables
                tr = page.rect
                tr.y1 = r.y0
                if not tr.is_empty:
                    text_rects.append(("text", tr, 0))
                text_rects.append(("table", r, idx))
                continue
            # read previous rectangle in final list: always a table!
            _, r0, idx0 = text_rects[-1]

            # check if a non-empty text rect is fitting in between tables
            tr = page.rect
            tr.y0 = r0.y1
            tr.y1 = r.y0
            if not tr.is_empty:  # empty if two tables overlap vertically!
                text_rects.append(("text", tr, 0))

            text_rects.append(("table", r, idx))

            # there may also be text below all tables
            if i == len(tab_rects) - 1:
                tr = page.rect
                tr.y0 = r.y1
                if not tr.is_empty:
                    text_rects.append(("text", tr, 0))

        if not text_rects:  # this will happen for table-free pages
            text_rects.append(("text", page.rect, 0))
        else:
            rtype, r, idx = text_rects[-1]
            if rtype == "table":
                tr = page.rect
                tr.y0 = r.y1
                if not tr.is_empty:
                    text_rects.append(("text", tr, 0))

        # we have all rectangles and can start outputting their contents
        for rtype, r, idx in text_rects:
            if rtype == "text":  # a text rectangle
                continue
            else:  # a table rect
                md_string += tabs[idx].to_markdown(clean=False)
                md_string += "</table>\n\n"

        md_string += "\n-----\n\n"

    return md_string


def extract_text_from_pdf(pdf_file: str) -> List[dict]:
    """
    Extracts text from a PDF file and returns a list of dictionaries containing the extracted text along with metadata.
    Tables are extracted as markdown but appended to the end of the text of each page instead of their original location. 
    Parameters:
        pdf_file (str): The path to the PDF file to extract text from.

    Returns:
        List[dict]: A list of dictionaries, each containing:
            - "text": The extracted text from the PDF page.
            - "source": The source of the text extraction (currently set to "pymupdf").
            - "file_directory": The directory of the PDF file.
            - "filename": The name of the PDF file without the extension.
            - "filetype": The extension of the PDF file.
            - "page_number": The page number corresponding to the extracted text.
    """

    all_res = []
    table_mds = to_markdown(pdf_file).split("\n-----\n\n") #this extracts only tables to markdown
    doc = fitz.open(pdf_file)
    page_no = 1
    for page, table_strs in zip(doc, table_mds):

        for tab in page.find_tables():
            # process the content of table 'tab'
            page.add_redact_annot(tab.bbox)  # wrap table in a redaction annotation

        page.apply_redactions()  # erase all table text
        page_text = page.get_text()

        if table_strs.strip():
            for table_str in table_strs.split("</table>\n\n"):
                page_text += "\n" + table_str + '\n'

        res = {
            "text": page_text.replace(" า", "ำ").replace(" ่า", "่ำ").replace(" ้า", "้ำ").replace(" ๊า", "๊ำ"),
            "source": "pymupdf",
            "file_directory": str(Path(pdf_file).parent),
            "filename": str(Path(pdf_file).stem),
            "filetype": Path(pdf_file).suffix,
            "page_number": page_no,
        }
        all_res.append(res)
        page_no+=1

    # TODO: add chunking
    return all_res