import os
import gradio as gr
from huggingface_hub import create_discussion
from huggingface_hub import hf_hub_download
from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes
from Crypto.Cipher import AES, PKCS1_OAEP
import base64

token = os.environ["CLAIM"]
repo = "Team8/dataset"

# <p style="text-align: center; color: black;"> 🎨 <span style='color: #6d071a;'>FindMy</span>Art</p>
description = """# <p style="text-align: center; color: black;"> 🎨 <span style='color: #6d071a;'>RemoveMy</span>Art</p>
<span>This is a space to opt-out your images from datasets. After you check that your image is in the dataset, 
fill an explanation for why you want it to be removed and specify if you want to encrypt the issue that will be opened on the dataset,
 by checking the Encrypt box.</span>"""


def encrypt_text(data):
    loc = hf_hub_download(repo_id="Team8/dataset", repo_type="dataset", filename="receiver.pem")
    data = data.encode("utf-8")
    file_out = open("encrypted_data.bin", "wb")

    recipient_key = RSA.import_key(open(loc).read())
    session_key = get_random_bytes(16)

    # Encrypt the session key with the public RSA key
    cipher_rsa = PKCS1_OAEP.new(recipient_key)
    enc_session_key = cipher_rsa.encrypt(session_key)

    # Encrypt the data with the AES session key
    cipher_aes = AES.new(session_key, AES.MODE_EAX)
    ciphertext, tag = cipher_aes.encrypt_and_digest(data)
    [ file_out.write(x) for x in (enc_session_key, cipher_aes.nonce, tag, ciphertext) ]
    file_out.close()
    with open("encrypted_data.bin", "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read())
    return encoded_string

def open_issue(explanation, email, encrypt):
    if encrypt:
        data = explanation + "\n User email: " + email
        explanation = f"Encrypted\n{encrypt_text(data)}"

    create_discussion(
        repo_id=repo,
        repo_type="dataset",
        title="[OPT-OUT REQUEST] Remove image from the dataset",
        description=explanation,
        token = token,
    )
    # to do add issue id to the link
    link = f"https://huggingface.co/datasets/{repo}/discussions"
    return f"Issue opened at {link}"

demo = gr.Blocks(
    css=".gradio-container {background-color: #fcf8ed; color:brown}"
)
with demo:
    with gr.Row():
        _, colum_2, _ = gr.Column(scale=1), gr.Column(scale=6), gr.Column(scale=1)
        with colum_2:
            gr.Markdown(value=description)
            explanation = gr.Textbox(lines=5, label="Please explain in a few lines why you want this image to be removed from the dataset.")
            email = gr.Textbox(lines=1, label="Fill your email if you want to be notified when the image is removed.")
            encrypt = gr.Checkbox(label="Encrypt my message")
            run = gr.Button("Open an issue on the dataset")
            output = gr.Textbox(lines=1, label="Opened issue link")
    event = run.click(open_issue, [explanation, email, encrypt],output, api_name="open_issue")

demo.launch()