import numpy as np
import pandas as pd
import requests
import json
import pickle
import io
import gradio as gr

from fhe_utils import FHEClient

from common import (
    ROOT_DIR,
    KEYS_PATH,
    DUNEDIN_PACE_DATA_PATH,
    SERVER_URL,
    SCALE,
    get_tmp_file_path
)

DUNEDIN_PACE_DATA = json.loads(open(DUNEDIN_PACE_DATA_PATH).read())


def get_client(user_id):
    return FHEClient(
        ROOT_DIR / "client.zip",
        scale=SCALE,
        dunedin_pace_data=DUNEDIN_PACE_DATA,
        key_dir=KEYS_PATH / f"{user_id}", 
    )

def generate_fhe_key():
    user_id = np.random.randint(0, 2**32)

    client = get_client(user_id)
    client.generate_private_and_evaluation_keys(force=True)
    evaluation_key = client.get_serialized_evaluation_keys()
    evaluation_key_path = get_tmp_file_path(user_id, "evaluation_key", "client")

    with evaluation_key_path.open("wb") as fw:
        fw.write(evaluation_key)

    return (user_id, True)


def save_to_file(user_id, file_name, data):
    file_path = get_tmp_file_path(user_id, file_name)
    with file_path.open("wb") as fw:
        fw.write(data)

def get_bytes(obj):
    obj_bytes = pickle.dumps(obj)
    return io.BytesIO(obj_bytes)

def get_pace_estimation(client, evaluation_key_path, betas_sample):
    betas_sample_es, weights_int_es, intercept_int_es = client.encrypt_serialize(betas_sample)

    files = [
        ("files", betas_sample_es),
        ("files", weights_int_es),
        ("files", intercept_int_es),
        ("files", open(evaluation_key_path, "rb")),
    ]
    resp = requests.post(
        f"{SERVER_URL}/pace_estimation",
        data={},
        files=files
    )
    pace_es = resp.content
    return client.deserialize_decrypt_post_process(pace_es)


def get_pace_estimations(user_id, betas_file):
    evaluation_key_path = get_tmp_file_path(user_id, "evaluation_key")

    try:
        beta_values = pd.read_csv(betas_file, delimiter="\t", index_col=0)
    except Exception:
        return (None,  gr.update(visible=False), "Non valid CSV")

    if len(beta_values.columns) == 0:
        return (None,  gr.update(visible=False), "[DATA ERROR]: No samples provided")

    for base_feature in DUNEDIN_PACE_DATA['probes']:
        if base_feature not in beta_values.index:
            return (None, None, f"[DATA ERROR]: Missing base probe '{base_feature}'")

    probes_mean = {
        DUNEDIN_PACE_DATA['normalization_probes'][i]: DUNEDIN_PACE_DATA['normalization_means'][i]
        for i in range(len(DUNEDIN_PACE_DATA['normalization_probes']))
    }
    missing_probes = [_ for _ in DUNEDIN_PACE_DATA['normalization_probes'] if _ not in beta_values.index]
    missing_probes_data = [[probes_mean[_]] * len(beta_values.columns) for _ in missing_probes]
    missing_betas = pd.DataFrame(missing_probes_data, index=missing_probes, columns=beta_values.columns)
    filtered_beta_values = pd.concat([beta_values, missing_betas]).loc[DUNEDIN_PACE_DATA['normalization_probes']]

    client = get_client(user_id)

    pace_data = [
        (column_name, get_pace_estimation(
            client, evaluation_key_path,
            filtered_beta_values[column_name].values
        ))
        for column_name in filtered_beta_values.columns
    ]

    return (pace_data, gr.update(visible=True), "Success")