from pathlib import Path
import gradio as gr
import pii_anonymizer

def upload_file(filepath):
    name = Path(filepath).name
    return [gr.UploadButton(visible=False), gr.DownloadButton(label=f"Download {name}", value=filepath, visible=True)]

def scrub_file(filepath, progress=gr.Progress()):
    # Temporary list to hold stripped lines
    stripped_lines = []
    pii_lines = []

    # Open the original file and strip each line
    with open(filepath, 'r', encoding='utf-8') as file:
        for line in file:
            stripped_lines.append(line.strip())
        progress(0, desc="Starting...")
        for line in progress.tqdm(stripped_lines):
            pii_lines.append(pii_anonymizer.anonymize_text_including_proper_nouns_and_addresses(line))

    # Overwrite the original file with pii replaced lines
    with open(filepath, 'w', encoding='utf-8') as file:
        for line in pii_lines:
            file.write(line + '\n')  # Add '\n' to keep the file structure

    # Extract the filename for the download button label
    name = Path(filepath).name

    # Return the Gradio interface elements
    return [gr.UploadButton(visible=False), gr.DownloadButton(label=f"Download {name}", value=filepath, visible=True)]


def download_file():
    return [gr.UploadButton(visible=True), gr.DownloadButton(visible=False)]

with gr.Blocks() as demo:
    gr.Markdown("Upload a file to create a PII stripped version, but it could take minutes before file will be ready for download.")
    #gr.Markdown("PII Anonymized:"+pii_anonymizer.anonymize_text_including_proper_nouns_and_addresses("Hello Kripa"))
    with gr.Row():
        u = gr.UploadButton("Upload a file", file_count="single", file_types=["text"])
        d = gr.DownloadButton("Download the file", visible=False)

    #u.upload(upload_file, u, [u, d])
    u.upload(scrub_file, u, [u, d])
    d.click(download_file, None, [u, d])


if __name__ == "__main__":
    demo.queue().launch()