import gradio as gr
from time import time
from pathlib import Path
from gradio_pdf import PDF
from pdf2image import convert_from_path
import shutil
import tempfile
from transformers import pipeline
import subprocess as sp

out_files = gr.State([])
FILE_TIMEOUT = 10 ** 3
MAX_FILES = 10

p = pipeline(
    "document-question-answering",
    model="TusharGoel/LiLT-Document-QA",
)

custom_css = """
<style>
button.primary {
    background-color: #2AD6B5 !important;
}
</style>
"""

def handle_files(cur_files):
    cur_time = cur_files[-1][0]
    deleted_indices = set()
    for other_idx, (other_time, other_file) in enumerate(cur_files[:-1]):
        if abs(cur_time - other_time) > FILE_TIMEOUT:
            shutil.rmtree(other_file.parent)
            deleted_indices.add(other_idx)
    cur_files = [cur_files[idx] for idx in range(len(cur_files)) if idx not in deleted_indices]
    
    if len(cur_files) > MAX_FILES:
        for _, other_file in cur_files[:-MAX_FILES]:
            shutil.rmtree(other_file.parent)
        cur_files = cur_files[-MAX_FILES:]
    return cur_files


def extract_text(pdf_file):
    """
    Generate a text rendering of a PDF file in the form of a list of lines.
    """
    args = ['pdftotext', '-layout', pdf_file, '-']
    cp = sp.run(
      args, stdout=sp.PIPE, stderr=sp.DEVNULL,
      check=True, text=True
    )
    return cp.stdout

# Function to process PDF and generate ZIP file
def process_pdf(pdf_file, cur_files):
    
    zip_output = Path(tempfile.mkdtemp()) / f'{Path(pdf_file).stem}'
    # zip_output.parent.mkdir()
    
    with tempfile.TemporaryDirectory() as path:
        pdf_output = path
        convert_from_path(pdf_file, output_folder=str(pdf_output))
    
        # Create a BytesIO object to store zip file in memory
        shutil.make_archive(zip_output, 'zip', pdf_output)
    
    zip_output = zip_output.with_suffix('.zip')
    
    cur_time = time()
    cur_files.append((cur_time, zip_output))
    cur_files = handle_files(cur_files)
    
    return str(zip_output), cur_files


def interact_with_pdf(doc, question):
    with tempfile.TemporaryDirectory() as path:
        images = convert_from_path(doc, output_folder=path)
        outputs = []
        for img in images:
            outputs += p(img, question)
    return sorted(outputs, key=lambda x: x["score"], reverse=True)[0]['answer']

'''
text_interface = gr.Interface(
    fn=extract_text,
    inputs=PDF(label="Загрузить PDF"),
    outputs=gr.Textbox(label="Полученный текст"),
    title="PDF в текст",
    description="Сервис вынимает содержащийся в PDF контейнере текст."
)

pdf_interface = gr.Interface(
    fn=process_pdf,
    inputs=[PDF(label="Загрузить PDF"), out_files],
    outputs=[gr.File(label="Скачать ZIP архив"), out_files],
    title="PDF в картинки",
    description="Переводит постранично файл в изображения, и предлагает бскачать в виде ZIP архива."
)

image_interface = gr.Interface(
    fn=interact_with_pdf,
    inputs=[
        PDF(label="Загрузить PDF"),
        gr.Textbox(label="Что хотите найти?")
    ],
    outputs=gr.Textbox(label="Возможный ответ"),
    title="Спроси PDF",
    description="Сервис через методы OCR сканирует загруженный файл для нахождения ответа на заданный вопрос."
)

# Create a tabbed interface
tabbed_interface = gr.TabbedInterface(
    [text_interface, pdf_interface, image_interface],
    title="Взаимодействие с PDF",
    tab_names=["В текст", "В картинки", "Задай вопрос"],
    # description="Choose a tab to perform the desired task.",
    css=custom_css,
    analytics_enabled=True
)
'''

theme = gr.themes.Default(primary_hue="blue").set(
    loader_color="#2AD6B5",
    slider_color="#2AD6B5",
    button_primary_background_fill="#2AD6B5",
    button_primary_text_color="#000000",
    button_primary_background_fill_hover="#20D0B0",
    button_primary_text_color_hover="#000000",
    button_primary_border_color="999999",
    button_primary_border_color_hover="333333"
)

with gr.Blocks(theme=theme) as tabbed_interface:
    with gr.Tab("В текст"):
        with gr.Row():
            with gr.Column():
                inp1 = PDF(label="Загрузить PDF")
                btn_sbm1 = gr.Button("Запустить", variant="primary")
            with gr.Column():
                out1 = gr.Textbox(label="Полученный текст")
        btn_sbm1.click(fn=extract_text, inputs=inp1, outputs=out1)
    with gr.Tab("В картинки"):
        with gr.Row():
            with gr.Column():
                inp2 = PDF(label="Загрузить PDF")
                btn_sbm2 = gr.Button("Запустить", variant="primary")
            with gr.Column():
                out2 = gr.File(label="Скачать ZIP архив")
        btn_sbm2.click(fn=process_pdf, inputs=[inp2, out_files], outputs=[out2, out_files])
    with gr.Tab("Задай вопрос"):
        with gr.Row():
            with gr.Column():
                inp3 = PDF(label="Загрузить PDF")
                inp4 = gr.Textbox(label="Что хотите найти?")
                btn_sbm3 = gr.Button("Запустить", variant="primary")
            with gr.Column():
                out3 = gr.Textbox(label="Возможный ответ")
        btn_sbm3.click(fn=interact_with_pdf, inputs=[inp3, inp4], outputs=out3)


tabbed_interface.launch(show_api=False, max_threads=8)