import argparse
import os
from typing import Literal

import gradio as gr
from PIL.Image import Image

from extractor import AzureExtractor
from utils import get_images_from_path
from dotenv import load_dotenv

load_dotenv()

parser = argparse.ArgumentParser()
parser.add_argument('-s', '--share', action='store_true', help='Create a public link')
args = parser.parse_args()

endpoint = os.environ['AZURE_ENDPOINT']
key = os.environ['AZURE_KEY']
extractor = AzureExtractor(endpoint=endpoint, key=key)


def handle_input(
    extract_mode: Literal['Content', 'Key Pairs'], file_input: str
) -> tuple[list[Image], dict[str, str]]:
    try:
        output_images = get_images_from_path(file_input)
        mode_mapping = {'Content': 'content', 'Key Pairs': 'key_pair'}
        result = extractor.extract(file_input, mode=mode_mapping[extract_mode])
        status = 'Success'
    except Exception as err:
        status = f'Error: {err}'
        output_images = None
        result = None

    return status, output_images, result


file_input = gr.File(file_types=['image', '.pdf'])
extract_mode = gr.Dropdown(
    ['Content', 'Key Pairs'], value='Content', label='Text extract mode'
)

status_box = gr.Textbox(label='Status')
image_gallery = gr.Gallery(label='Input Preview')
result = gr.Json(label='result')
demo = gr.Interface(
    fn=handle_input,
    inputs=[extract_mode, file_input],
    outputs=[status_box, image_gallery, result],
)
demo.launch(server_name='0.0.0.0', server_port=7860, share=args.share)