import fitz
from PIL import Image
from utils import *
from whoosh.analysis import Tokenizer, Token
import jieba
from whoosh.index import create_in
from whoosh.fields import *
from whoosh.qparser import QueryParser
import os
import shutil
# import tempfile

LOGO_WIDTH = 398
LOGO_HEIGHT = 137

ix = None
writer = None


class ChineseTokenizer(Tokenizer):
    def __call__(self, value, positions=False, chars=False,
                 keeporiginal=False, removestops=True,
                 start_pos=0, start_char=0, mode='', **kwargs):
        t = Token(positions, chars, removestops=removestops, mode=mode,
                  **kwargs)
        seglist = jieba.cut(value, cut_all=True)
        for w in seglist:
            t.original = t.text = w
            t.boost = 1.0
            if positions:
                t.pos = start_pos + value.find(w)
            if chars:
                t.startchar = start_char + value.find(w)
            if chars and positions:
                t.endchar = start_char + value.find(w) + len(w)
            yield t


def ChineseAnalyzer():
    return ChineseTokenizer()


def load_pdf(file, dpi=300, skip_page_front=0, skip_page_back=1, skip_block=5, lang='CN'):
    """
    Load pdf file, covert to image, description and index it
    :param lang:
    :param skip_block:
    :param skip_page_back:
    :param skip_page_front:
    :param dpi:
    :param file:
    :return:
    """

    if file.__contains__('\\gradio\\') or file.__contains__('/gradio/'):
        print('gradio file')
        doc = fitz.open(file)
    else:
        print('local file')
        doc = fitz.open('using_pdfs/' + file)

    # load pages
    pages = []
    for i in range(doc.page_count):
        page = doc.load_page(i)
        pages.append(page)

    # increase dpi to 300
    dpi = int(dpi)
    scale = dpi / 72  # default dpi of pdf is 72
    matrix = fitz.Matrix(scale, scale)
    skip_block = int(skip_block)

    base_name = os.path.basename(file).split('.')[0]
    path_name = f'images/{base_name}'
    if os.path.exists(path_name):
        shutil.rmtree(path_name)
    os.mkdir(path_name)

    temp_image_dir = path_name
    # temp_image_dir = tempfile.mkdtemp(prefix='images_')

    for page in pages[int(skip_page_front):-int(skip_page_back)]:  # skip final page

        # part1: get image with description in png-pdf
        p1dict = page.get_text('dict')
        blocks = p1dict['blocks']
        page_pix = page.get_pixmap(matrix=matrix, dpi=dpi)
        page_im = Image.frombytes("RGB", (page_pix.width, page_pix.height), page_pix.samples)

        saved = []  # need to remove if inner a svg image
        for i, block in enumerate(blocks[int(skip_block):]):  # head and tail of pages should be ignore
            if 'image' in block:
                # try:
                    bbox = block['bbox']
                    # skip image that width=398 and hight=137 -> Typically LOGO
                    if (bbox[2] - bbox[0])*scale - LOGO_WIDTH <= 10 and (bbox[3] - bbox[1])*scale - LOGO_HEIGHT <= 10:
                        continue
                    # Scale the bbox coordinates
                    cropped = page_im.crop([int(i * scale) for i in bbox])
                    number = block['number']

                    file_name = temp_image_dir + f'/{base_name}_imgbmp_{page.number}_{number}'
                    image_name = file_name + '.png'
                    # print(image_name)
                    cropped.save(image_name)
                    # # Handle text extraction around the image
                    text_content = get_text_around_image(blocks[skip_block:], i, lang)
                    title = get_title_of_image(blocks[skip_block:], i, lang)
                    # print(text_content[:30])
                    # print(title)
                    with open(f'{file_name}.txt', 'w', encoding='utf-8') as text_file:
                        text_file.write(title + '\n' + text_content.replace('\n', ' ')+ f'\nbase name:{base_name}')

                    saved.append((file_name, [int(i * scale) for i in bbox]))
                # except:
                #     pass

        # part2: get image with description in svg-pdf
        svg = page.get_svg_image(matrix=fitz.Identity)
        image_clips, svg_blocks = parse_page_svg(svg, page.number)
        for clip in image_clips:

            transform = []
            for item in clip[0]:
                # print(item, type(item))
                if item[0] == '.':
                    transform.append(float('0' + item))
                elif item[0] == '-':
                    transform.append(float('-0' + item[1:]))
                else:
                    transform.append(float(item))
            d = clip[1]
            page_id = clip[2]
            block_id = clip[3]
            matches = re.findall(r'H(\d+\.?\d*)V(\d+\.?\d*)', d)
            float_values = [float(value) for value in matches[0]]
            box_width = float_values[0]
            box_height = float_values[1]
            width_scale = transform[0]
            height_scale = transform[3]
            width_move = transform[4]
            height_move = transform[5]
            x1 = width_move * scale
            y1 = height_move * scale
            # x1=347*scale
            # y1=587*scale
            x2 = x1 + box_width * width_scale * scale
            y2 = y1 + box_height * height_scale * scale

            if y1 > y2:
                y1, y2 = y2, y1

            # print(x1, y1, x2, y2)
            # 3. 截取并保存图像

            # check images in saved, if in or similar, delete it from file system
            for i, (file_name, bbox) in enumerate(saved):
                if (abs(bbox[0] - x1) < 10\
                        and abs(bbox[1] - y1) < 10\
                        and abs(bbox[2] - x2) < 10\
                        and abs(bbox[3] - y2) < 10) or \
                        (bbox[0]>x1-10 and bbox[1]>y1-10 and bbox[2]<x2+10 and bbox[3]<y2+10):
                    os.remove(file_name + '.png')
                    os.remove(file_name + '.txt')
                    saved.pop(i)
                    break

            cropped_img = page_im.crop((int(x1), int(y1), int(x2), int(y2)))
            file_name = temp_image_dir + f'/{base_name}_imgsvg_{page.number}_{block_id}'
            image_name = file_name + '.png'
            cropped_img.save(image_name)

            # search title and text
            text_content = get_svg_text_around_image(svg_blocks, block_id, lang)
            title = get_svg_title_around_image(svg_blocks, block_id, lang)
            with open(f'{file_name}.txt', 'w', encoding='utf-8') as text_file:
                text_file.write(title + '\n' + text_content.replace('\n', ' ') + f'\nbase name:{base_name}')

    print(temp_image_dir)
    return temp_image_dir


def build_index(file, tmp_dir, lang='CN'):
    # Define the schema for the index
    if lang == 'CN':
        schema = Schema(file_name=ID(stored=True), content=TEXT(analyzer=ChineseAnalyzer(), stored=True))
    else:
        schema = Schema(file_name=ID(stored=True), content=TEXT(stored=True))

    base_name = os.path.basename(file).split('.')[0]
    path_name = f'{base_name}'
    # index_path = 'indexes/' + path_name + '_index_dir'
    index_path = 'indexes/'
    # Create an index in a directory
    # if os.path.exists(index_path):
    #     shutil.rmtree(index_path)
    # os.mkdir(index_path)
    temp_index_dir = index_path
    # temp_index_dir = tempfile.mkdtemp(prefix='index_')

    global ix
    if ix is None:
        ix = create_in(temp_index_dir, schema)
    global writer
    if writer is None:
        writer = ix.writer()

    # Add documents to the index
    # base_name = os.path.basename(file).split('.')[0]
    # image_path = f'images{base_name}'
    # writer = ix.writer()
    for file in os.listdir(tmp_dir):
        if file.endswith('.txt'):
            file_path = os.path.join(tmp_dir, file)
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            writer.add_document(file_name=file[:-4], content=content)
            print('==========')
            print(content)
            print("==========")

    writer.commit()
    return ix, temp_index_dir


def search(ix, query, lang='CN', k=10):

    # Tokenize the query string and join tokens with OR operator
    if lang == 'CN':
        query_tokens = jieba.cut(query, cut_all=True)
    else:
        query_tokens = query.split()
    or_query = " OR ".join(query_tokens)

    parser = QueryParser("content", ix.schema)
    myquery = parser.parse(or_query)

    with ix.searcher() as searcher:
        results = searcher.search(myquery, limit=k)

        # Extract and return the file names and descriptions of the top-k hits
        results_list = [(hit['file_name'], hit['content'], hit.score) for hit in results]

    return results_list


def return_image(file, results_list, tmp_dir):
    # base_name = os.path.basename(file).split('.')[0]
    # path_name = f'images{base_name}'
    titles = []
    images = []
    for result in results_list:
        title = result[1].split('\n')[0].split(':')[-1]
        titles.append(title)
        images.append(Image.open(tmp_dir + '/' + result[0] + '.png'))
    return titles[0], images[0]


# file = 'CA-IS372x-datasheet_cn.pdf'
# file = 'CA-IS3086 datasheet_cn.pdf'
# temp_image_dir = load_pdf(file, lang='CN')
# ix, temp_index_dir = build_index(file, temp_image_dir)
# results_list = search(ix, "波形", lang='CN', k=10)
# ret_img = return_image(file, results_list, temp_image_dir)
# print('title: ' + ret_img[0])
# ret_img[1].show()

# print(os.listdir('using_pdfs'))

# import tqdm
# for file in tqdm.tqdm(os.listdir('using_pdfs')):
#     tmd_dir = load_pdf(file)
#     ix, tmp_index_dir = build_index('using_pdfs/' + file, tmd_dir)
# #
# writer.commit()

# from whoosh.index import open_dir
# search_ix = open_dir('indexes')
# query = "IF-428x接收端阈值"
# results = search(search_ix, query, lang='CN', k=10)
# for result in results:
#     print(result)
#
# from PIL import Image
#
# for result in results:
#     image_name = result[0]
#     base_name = image_name.split('_img')[0]
#     img = Image.open('images/' + base_name + '/' + image_name + '.png')
#     image_title = result[1].split('\n')[0].split(':')[1]
#     img.show(title=image_title)