import gradio as gr
import time
import os
from pathlib import Path

import subprocess
from concrete.ml.deployment import FHEModelClient
from requests import head
import numpy
import os
from pathlib import Path
import requests
import json
import base64
import subprocess
import shutil
import time
import pandas as pd
import pickle
import numpy as np

# This repository's directory
REPO_DIR = Path(__file__).parent
subprocess.Popen(["uvicorn", "server:app"], cwd=REPO_DIR)

# if not exists, create a directory for the FHE keys called .fhe_keys
if not os.path.exists(".fhe_keys"):
    os.mkdir(".fhe_keys")
# if not exists, create a directory for the tmp files called tmp
if not os.path.exists("tmp"):
    os.mkdir("tmp")


# Wait 4 sec for the server to start
time.sleep(4)


# Encrypted data limit for the browser to display
# (encrypted data is too large to display in the browser)
ENCRYPTED_DATA_BROWSER_LIMIT = 500
N_USER_KEY_STORED = 20


def clean_tmp_directory():
    # Allow 20 user keys to be stored.
    # Once that limitation is reached, deleted the oldest.
    path_sub_directories = sorted(
        [f for f in Path(".fhe_keys/").iterdir() if f.is_dir()], key=os.path.getmtime
    )

    user_ids = []
    if len(path_sub_directories) > N_USER_KEY_STORED:
        n_files_to_delete = len(path_sub_directories) - N_USER_KEY_STORED
        for p in path_sub_directories[:n_files_to_delete]:
            user_ids.append(p.name)
            shutil.rmtree(p)

    list_files_tmp = Path("tmp/").iterdir()
    # Delete all files related to user_id
    for file in list_files_tmp:
        for user_id in user_ids:
            if file.name.endswith(f"{user_id}.npy"):
                file.unlink()


def keygen():
    # Clean tmp directory if needed
    clean_tmp_directory()

    print("Initializing FHEModelClient...")
    # Let's create a user_id
    user_id = numpy.random.randint(0, 2**32)
    fhe_api = FHEModelClient(f"deployment/deployment_{task}", f".fhe_keys/{user_id}")
    fhe_api.load()

    # Generate a fresh key
    fhe_api.generate_private_and_evaluation_keys(force=True)
    evaluation_key = fhe_api.get_serialized_evaluation_keys()

    numpy.save(f"tmp/tmp_evaluation_key_{user_id}.npy", evaluation_key)

    return [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], user_id]


def encode_quantize_encrypt(test_file, user_id):

    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
    fhe_api.load()
    from PE_main import extract_infos

    features = pickle.loads(open(os.path.join("features.pkl"), "rb").read())
    encodings = extract_infos(test_file)
    encodings = list(map(lambda x: encodings[x], features))

    quantized_encodings = fhe_api.model.quantize_input(encodings).astype(numpy.uint8)
    encrypted_quantized_encoding = fhe_api.quantize_encrypt_serialize(encodings)

    # Save encrypted_quantized_encoding in a file, since too large to pass through regular Gradio
    # buttons, https://github.com/gradio-app/gradio/issues/1877
    numpy.save(
        f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy",
        encrypted_quantized_encoding,
    )

    # Compute size
    encrypted_quantized_encoding_shorten = list(encrypted_quantized_encoding)[
        :ENCRYPTED_DATA_BROWSER_LIMIT
    ]
    encrypted_quantized_encoding_shorten_hex = "".join(
        f"{i:02x}" for i in encrypted_quantized_encoding_shorten
    )
    return (
        encodings[0],
        quantized_encodings[0],
        encrypted_quantized_encoding_shorten_hex,
    )


def run_fhe(user_id):
    encoded_data_path = Path(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy")
    encrypted_quantized_encoding = numpy.load(encoded_data_path)

    # Read evaluation_key from the file
    evaluation_key = numpy.load(f"tmp/tmp_evaluation_key_{user_id}.npy")

    # Use base64 to encode the encodings and evaluation key
    encrypted_quantized_encoding = base64.b64encode(
        encrypted_quantized_encoding
    ).decode()
    encoded_evaluation_key = base64.b64encode(evaluation_key).decode()

    query = {}
    query["evaluation_key"] = encoded_evaluation_key
    query["encrypted_encoding"] = encrypted_quantized_encoding
    headers = {"Content-type": "application/json"}


    response = requests.post(
        "http://localhost:8000/predict",
        data=json.dumps(query),
        headers=headers,
    )


    encrypted_prediction = base64.b64decode(response.json()["encrypted_prediction"])

    numpy.save(f"tmp/tmp_encrypted_prediction_{user_id}.npy", encrypted_prediction)
    encrypted_prediction_shorten = list(encrypted_prediction)[
        :ENCRYPTED_DATA_BROWSER_LIMIT
    ]
    encrypted_prediction_shorten_hex = "".join(
        f"{i:02x}" for i in encrypted_prediction_shorten
    )


def decrypt_prediction(user_id):
    encoded_data_path = Path(f"tmp/tmp_encrypted_prediction_{user_id}.npy")

    # Read encrypted_prediction from the file

    encrypted_prediction = numpy.load(encoded_data_path).tobytes()

    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
    fhe_api.load()

    # We need to retrieve the private key that matches the client specs (see issue #18)
    fhe_api.generate_private_and_evaluation_keys(force=False)

    predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)


def update(name):
    return f"Welcome to Gradio, {name}!"


if __name__ == "__main__":
    app = gr.Interface(
        [
            keygen,
            encode_quantize_encrypt,
            run_fhe,
            decrypt_prediction,
        ],
        [
            gr.inputs.Textbox(label="Task", default="malware"),
            gr.inputs.File(label="Test File"),
            gr.inputs.Textbox(label="User ID"),
        ],
        [
            gr.outputs.Textbox(label="Evaluation Key"),
            gr.outputs.Textbox(label="Encodings"),
            gr.outputs.Textbox(label="Encrypted Quantized Encoding"),
            gr.outputs.Textbox(label="Encrypted Prediction"),
        ],
        title="FHE Model",
        description="This is a FHE Model",
    )
    app.launch()