import gradio as gr
import os
from pii_transform.api.e2e import PiiTextProcessor
from pii_extract.defs import FMT_CONFIG_PLUGIN

examples = []
with open("examples.txt", "r") as f:
    examples = f.readlines()
examples_truncated = [example[:50] + "..." for example in examples]
language_choices = {
    "English": "en",
    "Italian": "it",
    "Spanish": "es",
    "Portugese": "pt",
    "Deutsche": "de",
    "French": "fr",
}
language_code = "en"
cache_dir = "/home/user/app/cache"
os.makedirs(cache_dir, exist_ok=True)
if os.path.isdir(cache_dir):
    gr.Info("Cache directory created at "+cache_dir)
else:
    gr.Warning("Cache directory creation error")

def change_language(language_selection):
    global language_code
    language_code = language_choices[language_selection]
    gr.Info(f"{language_selection} selected")


def process(text, policy):
    # Create the object, defining the language to use and the policy
    # Further customization is possible by providing a config
    if text == "":
        print("Empty text field")
        gr.Warning("No text present")
        return ""

    # Custom config to prevent loading of the Presidio plugin
    # config = {FMT_CONFIG_PLUGIN: {"piisa-detectors-presidio": {"load": False}}}
    proc = PiiTextProcessor(
        lang=language_code, default_policy=policy, config="config.json"
    )

    # Process a text buffer and get the transformed buffer
    outbuf = proc(text)
    return outbuf


def get_full_example(idx):
    return examples[idx]


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=0, min_width=75):
            logo = gr.Image("image.jpeg", show_label=False, show_download_button=False)
        with gr.Column():
            pass
        with gr.Column(scale=0, min_width=200):
            lang_picker = gr.Dropdown(
                choices=list(language_choices.keys()),
                label="Select Language",
                value=list(language_choices.keys())[0],
                type="value",
            )
            lang_picker.select(change_language, inputs=lang_picker, outputs=None)
    with gr.Row():
        with gr.Column(scale=2, min_width=400):
            text_original = gr.Textbox(
                label="Original Text",
                lines=10,
                placeholder="Enter the text you would like to analyze, or select from one of the examples below",
            )
        with gr.Column(scale=0, min_width=25):
            pass
        with gr.Column(scale=0, min_width=100):
            for i in range(3):
                with gr.Row():
                    pass
            redact_btn = gr.Button(value="Redact", variant="primary", size="sm")
            anonymize_btn = gr.Button(value="Anonymize", variant="primary", size="sm")
            placeholder_btn = gr.Button(
                value="Placeholder", variant="primary", size="sm"
            )
        with gr.Column(scale=0, min_width=25):
            pass
        with gr.Column(
            scale=2,
            min_width=400,
        ):
            text_redacted = gr.TextArea(
                label="Transformed Text",
                lines=10,
                show_copy_button=True,
                interactive=False,
            )

        redact_btn.click(
            fn=process,
            inputs=[
                text_original,
                gr.Text(value="redact", visible=False),
            ],
            outputs=text_redacted,
        )
        anonymize_btn.click(
            fn=process,
            inputs=[
                text_original,
                gr.Text(value="synthetic", visible=False),
            ],
            outputs=text_redacted,
        )
        placeholder_btn.click(
            fn=process,
            inputs=[
                text_original,
                gr.Text(value="label", visible=False),
            ],
            outputs=text_redacted,
        )
    with gr.Row():
        example_selector = gr.Dropdown(
            examples_truncated, type="index", label="Examples"
        )
        example_selector.select(
            get_full_example, inputs=example_selector, outputs=text_original
        )

demo.queue().launch()