import spaces
import gradio as gr
from marker.markdown_extractor import MarkdownExtractorConfig, MarkdownExtractor
from pdf.pdf_extractor import PDFExtractorConfig, PDFExtractor
from indexify_extractor_sdk import Content

markdown_extractor = MarkdownExtractor()
pdf_extractor = PDFExtractor()

@spaces.GPU
def use_marker(pdf_filepath):
	if pdf_filepath is None:
		raise gr.Error("Please provide some input PDF: upload a PDF file")
	with open(pdf_filepath, "rb") as f:
		pdf_data = f.read()
	content = Content(content_type="application/pdf", data=pdf_data)
	config = MarkdownExtractorConfig(batch_multiplier=2)
	result = markdown_extractor.extract(content, config)
	return result

with gr.Blocks(title="PDF data extraction with Marker & Indexify") as marker_demo:
	gr.HTML("<h1 style='text-align: center'>PDF data extraction with Marker & <a href='https://getindexify.ai/'>Indexify</a></h1>")
	gr.HTML("<p style='text-align: center'>Indexify is a scalable realtime and continuous indexing and structured extraction engine for unstructured data to build generative AI applications</p>")
	gr.HTML("<h3 style='text-align: center'>If you like this demo, please ⭐ Star us on <a href='https://github.com/tensorlakeai/indexify' target='_blank'>GitHub</a>!</h3>")
	gr.HTML("<h4 style='text-align: center'>Here's an example notebook that demonstrates how to build a continuous <a href='https://github.com/tensorlakeai/indexify/blob/main/docs/docs/examples/efficient_rag.ipynb' target='_blank'>extraction pipeline</a> with Indexify</h4>")

	with gr.Row():
		with gr.Column():
			gr.HTML(
				"<p><b>Step 1:</b> Upload a PDF file from local storage.</p>"
				"<p style='color: #A0A0A0;'>Use this demo for single PDF file only. "
				"You can extract from PDF files continuously and try various other extractors locally with "
				"<a href='https://getindexify.ai/'>Indexify</a>.</p>"
			)
			pdf_file_1 = gr.File(type="filepath")
		with gr.Column():
			gr.HTML("<p><b>Step 2:</b> Run the extractor.</p>")
			go_button = gr.Button(value="Run extractor", variant="primary")
			model_output_text_box_1 = gr.Textbox(label="Extractor Output", elem_id="model_output_text_box_1")

	with gr.Row():
		gr.HTML("<p style='text-align: center'>Developed with 🫶 by <a href='https://getindexify.ai/' target='_blank'>Indexify</a> | a <a href='https://www.tensorlake.ai/' target='_blank'>Tensorlake</a> product</p>")

	go_button.click(fn=use_marker, inputs=[pdf_file_1], outputs=[model_output_text_box_1])

@spaces.GPU
def use_pdf_extractor(pdf_filepath):
	if pdf_filepath is None:
		raise gr.Error("Please provide some input PDF: upload a PDF file")
	with open(pdf_filepath, "rb") as f:
		pdf_data = f.read()
	content = Content(content_type="application/pdf", data=pdf_data)
	config = PDFExtractorConfig(output_types=["text", "table"])
	result = pdf_extractor.extract(content, config)
	return result

with gr.Blocks(title="PDF data extraction with PDF Extractor & Indexify") as pdf_demo:
	gr.HTML("<h1 style='text-align: center'>PDF data extraction with PDF Extractor & <a href='https://getindexify.ai/'>Indexify</a></h1>")
	gr.HTML("<p style='text-align: center'>Indexify is a scalable realtime and continuous indexing and structured extraction engine for unstructured data to build generative AI applications</p>")
	gr.HTML("<h3 style='text-align: center'>If you like this demo, please ⭐ Star us on <a href='https://github.com/tensorlakeai/indexify' target='_blank'>GitHub</a>!</h3>")
	gr.HTML("<h4 style='text-align: center'>Here's an example notebook that demonstrates how to build a continuous <a href='https://github.com/tensorlakeai/indexify/blob/main/docs/docs/examples/SEC_10_K_docs.ipynb' target='_blank'>extraction pipeline</a> with Indexify</h4>")

	with gr.Row():
		with gr.Column():
			gr.HTML(
				"<p><b>Step 1:</b> Upload a PDF file from local storage.</p>"
				"<p style='color: #A0A0A0;'>Use this demo for single PDF file only. "
				"You can extract from PDF files continuously and try various other extractors locally with "
				"<a href='https://getindexify.ai/'>Indexify</a>.</p>"
			)
			pdf_file_2 = gr.File(type="filepath")
		with gr.Column():
			gr.HTML("<p><b>Step 2:</b> Run the extractor.</p>")
			go_button = gr.Button(value="Run extractor", variant="primary")
			model_output_text_box_2 = gr.Textbox(label="Extractor Output", elem_id="model_output_text_box_2")

	with gr.Row():
		gr.HTML("<p style='text-align: center'>Developed with 🫶 by <a href='https://getindexify.ai/' target='_blank'>Indexify</a> | a <a href='https://www.tensorlake.ai/' target='_blank'>Tensorlake</a> product</p>")

	go_button.click(fn=use_pdf_extractor, inputs=[pdf_file_2], outputs=[model_output_text_box_2])

demo = gr.TabbedInterface([marker_demo, pdf_demo], ["Marker Extractor", "PDF Extractor"], theme=gr.themes.Soft())

demo.queue()
demo.launch()