import streamlit as st
import streamlit.components.v1 as components
import random
from pathlib import Path
import json
import os
from huggingface_hub import upload_file
from datetime import datetime
from cryptography.fernet import Fernet as F
from typing import Union

HF_TOKEN = os.environ.get("HF_TOKEN")
print("is none?", HF_TOKEN is None)
FERNET_KEY = os.environ.get("FERNET_KEY")
print("FERNET_KEY is none?", FERNET_KEY is None)

# https://stackoverflow.com/questions/61607367/how-to-encrypt-json-in-python

def encrypt_message(key: str, msg: Union[str, bytes]) -> str:
    encoded_key = key if isinstance(key, bytes) else key.encode()
    handler = F(encoded_key)
    encoded_msg = msg if isinstance(msg, bytes) else msg.encode()
    treatment = handler.encrypt(encoded_msg)
    return str(treatment, 'utf-8')

def decrypt_message(key: str, msg: Union[str, bytes]) -> str:
    encoded_key = key if isinstance(key, bytes) else key.encode()
    handler = F(encoded_key)
    encoded_msg = msg if isinstance(msg, bytes) else msg.encode()
    treatment = handler.decrypt(encoded_msg)
    return str(treatment, 'utf-8')

user = st.sidebar.text_input("User", "your.email@provider.com")

paths_to_html = random.shuffle(
    [Path(filename) for filename in Path("./html_files/").glob("*.html")]
)

label_to_number = {"Human": 0, "AI": 1}

if "paths_to_html" not in st.session_state:
    st.session_state.paths_to_html = [
        Path(filename) for filename in Path("./html_files/").glob("*.html")
    ]
    random.shuffle(st.session_state.paths_to_html)

if "html_idx" not in st.session_state:
    st.session_state.html_idx = 0

if "html_key" not in st.session_state:
    st.session_state.html_key = ""

if "html_ratings" not in st.session_state:
    st.session_state.html_ratings = {}

# Display the image
path_to_html = st.session_state.paths_to_html[st.session_state.html_idx]
with open(path_to_html, "r") as f:
    html_data = f.read()
st.components.v1.html(html_data, height=2100)

if st.session_state["html_key"] in st.session_state.html_ratings:
    rating = st.session_state.html_ratings[st.session_state["html_key"]]
else:
    rating = ""

if rating := st.sidebar.radio(
    "Written by a human or generated by AI ?",
    ["Human", "AI"],
    horizontal=True,
    index=None if rating == "" else label_to_number[rating],
    key=path_to_html,
):
    key = st.session_state.paths_to_html[st.session_state.html_idx].stem
    st.session_state.html_key = key
    st.session_state.html_ratings[key] = rating
    st.sidebar.info("Submitted!")

if st.sidebar.button("Next email", key="next"):
    idx = st.session_state.html_idx
    idx = (idx + 1) % len(st.session_state.paths_to_html)
    st.session_state.html_idx = idx
    st.session_state.html_key = st.session_state.paths_to_html[
        st.session_state.html_idx
    ].stem
    st.rerun()

elif st.sidebar.button("Previous email", key="previous"):
    idx = st.session_state.html_idx
    idx = (idx - 1) % len(st.session_state.paths_to_html)
    st.session_state.html_idx = idx
    st.session_state.html_key = st.session_state.paths_to_html[
        st.session_state.html_idx
    ].stem
    st.rerun()
else:
    if st.sidebar.button("Quit", key="quit"):
        # send email with data or append to jsonl file
        if user == "your.email@provider.com":
            st.sidebar.error(f"Please provide your email adress")
        else:
            encrypted_user = encrypt_message(FERNET_KEY, user)
            assert user == decrypt_message(FERNET_KEY, encrypted_user), "decoded is different from original"
            with open("rate_results.jsonl", "a", encoding="utf-8") as jsonl:
                rating_data = {"user": encrypted_user, "ratings": st.session_state.html_ratings, "timestamp": datetime.now().isoformat()}
                jsonl.write(json.dumps(rating_data, ensure_ascii=False))
                jsonl.write("\n")
            upload_file(
                path_or_fileobj="rate_results.jsonl",
                path_in_repo="rate_results.jsonl",
                repo_id="simonn8/email_rater",
                repo_type="space",
                token=HF_TOKEN,
                commit_message=f"ratings by {encrypted_user}"
            )
            st.sidebar.success(
                f"Thank you for rating {len(st.session_state.html_ratings)} emails. Registered for user {user}. You can close the webpage."
            )



# st.session_state