import streamlit as st
import json
import os
from cryptography.fernet import Fernet

# Initialize session state
if 'username' not in st.session_state:
    st.session_state.username = ''
if 'scratch_pad' not in st.session_state:
    st.session_state.scratch_pad = ''

# Load or create user history file
history_file = 'user_history.json'
if not os.path.exists(history_file):
    with open(history_file, 'w') as file:
        json.dump({}, file)

# Load user history
with open(history_file, 'r') as file:
    user_history = json.load(file)

# Generate encryption key
key = Fernet.generate_key()
fernet = Fernet(key)

def save_session(username, password):
    st.session_state.username = username
    st.session_state.scratch_pad = st.session_state.scratch_pad
    encrypted_password = fernet.encrypt(password.encode()).decode()
    user_history[username] = encrypted_password
    with open(history_file, 'w') as file:
        json.dump(user_history, file)

def register_user(username, password):
    if username not in user_history:
        save_session(username, password)
        st.success("User registered successfully!")
    else:
        st.warning("Username already exists. Please choose a different username.")

def login_user(username, password):
    if username in user_history:
        encrypted_password = user_history[username]
        decrypted_password = fernet.decrypt(encrypted_password.encode()).decode()
        if password == decrypted_password:
            save_session(username, password)
            st.success("Logged in successfully!")
        else:
            st.error("Invalid password. Please try again.")
    else:
        st.warning("Username not found. Please register first.")

def reset_password(username):
    if username in user_history:
        new_password = Fernet.generate_password(8).decode()
        save_session(username, new_password)
        st.success(f"Password reset successful. Your new password is: {new_password}")
    else:
        st.warning("Username not found.")

st.title("Welcome to My App")

if st.session_state.username:
    st.subheader(f"Welcome back, {st.session_state.username}!")
    scratch_pad = st.text_area("Scratch Pad", value=st.session_state.scratch_pad, height=200)
    if st.button("💾 Save"):
        save_session(st.session_state.username, '')
    if st.button("🔑 Reset Password"):
        reset_password(st.session_state.username)
else:
    username = st.text_input("Username:")
    password = st.text_input("Password:", type="password")
    col1, col2 = st.columns(2)
    with col1:
        if st.button("🔑 Login"):
            login_user(username, password)
    with col2:
        if st.button("📝 Register"):
            register_user(username, password)