def hash_passkey(passkey):
    return hashlib.sha256(passkey.encode()).hexdigest()

def generate_key_from_passkey(passkey):
    hashed = hashlib.sha256(passkey.encode()).digest()
    return base64.urlsafe_b64encode(hashed[:32])

def encrypt_data(text, passkey):
    key = generate_key_from_passkey(passkey)
    cipher = Fernet(key)
    return cipher.encrypt(text.encode()).decode()

def decrypt_data(encrypted_text, passkey, data_id):
    try:
        hashed_passkey = hash_passkey(passkey)
        if (
            data_id in streamlit.session_state.stored_data and
            streamlit.session_state.stored_data[data_id]["passkey"] == hashed_passkey
        ):
            key = generate_key_from_passkey(passkey)
            cipher = Fernet(key)
            decrypted = cipher.decrypt(encrypted_text.encode()).decode()
            streamlit.session_state.failed_attempts = 0
            return decrypted
        else:
            streamlit.session_state.failed_attempts += 1
            streamlit.session_state.last_attempt_time = time.time()
            return None
    except Exception:
        streamlit.session_state.failed_attempts += 1
        streamlit.session_state.last_attempt_time = time.time()
        return None

def reset_failed_attempts():
    streamlit.session_state.failed_attempts = 0

def change_page(page):
    streamlit.session_state.current_page = page

# ---------------------- Main App ----------------------

def main():
    # Initialize session state
    if "failed_attempts" not in streamlit.session_state:
        streamlit.session_state.failed_attempts = 0
    if "stored_data" not in streamlit.session_state:
        streamlit.session_state.stored_data = {}
    if "current_page" not in streamlit.session_state:
        streamlit.session_state.current_page = "Home"
    if "last_attempt_time" not in streamlit.session_state:
        streamlit.session_state.last_attempt_time = 0

    streamlit.title("🔒 Secure Data Encryption System")

    # Navigation
    menu = ["Home", "Store Data", "Retrieve Data", "Login"]
    choice = streamlit.sidebar.selectbox("Navigation", menu, index=menu.index(streamlit.session_state.current_page))
    streamlit.session_state.current_page = choice

    # Handle lockout
    if streamlit.session_state.failed_attempts >= 3:
        streamlit.session_state.current_page = "Login"
        streamlit.warning("🔒 Too many failed attempts! Reauthorization required.")

    # ---------------------- Pages ----------------------

    # Home Page
    if streamlit.session_state.current_page == "Home":
        streamlit.subheader("🏠 Welcome to the Secure Data System")
        streamlit.write("Use this app to **securely store and retrieve data** using unique passkeys.")

        col1, col2 = streamlit.columns(2)
        with col1:
            if streamlit.button("Store New Data", use_container_width=True):
                change_page("Store Data")
        with col2:
            if streamlit.button("Retrieve Data", use_container_width=True):
                change_page("Retrieve Data")

        streamlit.info(f"Currently storing {len(streamlit.session_state.stored_data)} encrypted data entries.")

    # Store Data Page
    elif streamlit.session_state.current_page == "Store Data":
        streamlit.subheader("📂 Store Data Securely")
        user_data = streamlit.text_area("Enter Data:")
        passkey = streamlit.text_input("Enter Passkey:", type="password")
        confirm_passkey = streamlit.text_input("Confirm Passkey:", type="password")

        if streamlit.button("Encrypt & Save"):
            if user_data and passkey and confirm_passkey:
                if passkey != confirm_passkey:
                    streamlit.error("⚠️ Passkeys do not match!")
                else:
                    data_id = str(uuid.uuid4())
                    hashed_passkey = hash_passkey(passkey)
                    encrypted_text = encrypt_data(user_data, passkey)

                    streamlit.session_state.stored_data[data_id] = {
                        "encrypted_text": encrypted_text,
                        "passkey": hashed_passkey
                    }

                    streamlit.success("✅ Data stored securely!")
                    streamlit.code(data_id, language="text")
                    streamlit.info("⚠️ Save this Data ID! You'll need it to retrieve your data.")
            else:
                streamlit.error("⚠️ All fields are required!")

    # Retrieve Data Page
    elif streamlit.session_state.current_page == "Retrieve Data":
        streamlit.subheader("🔍 Retrieve Your Data")
        attempts_remaining = 3 - streamlit.session_state.failed_attempts
        streamlit.info(f"Attempts remaining: {attempts_remaining}")

        data_id = streamlit.text_input("Enter Data ID:")
        passkey = streamlit.text_input("Enter Passkey:", type="password")

        if streamlit.button("Decrypt"):
            if data_id and passkey:
                if data_id in streamlit.session_state.stored_data:
                    encrypted_text = streamlit.session_state.stored_data[data_id]["encrypted_text"]
                    decrypted_text = decrypt_data(encrypted_text, passkey, data_id)

                    if decrypted_text:
                        streamlit.success("✅ Decryption successful!")
                        streamlit.markdown("### Your Decrypted Data:")
                        streamlit.code(decrypted_text, language="text")
                    else:
                        streamlit.error(f"❌ Incorrect passkey! Attempts remaining: {3 - streamlit.session_state.failed_attempts}")
                else:
                    streamlit.error("❌ Data ID not found!")

                if streamlit.session_state.failed_attempts >= 3:
                    streamlit.warning("🔒 Too many failed attempts! Redirecting to Login Page.")
                    streamlit.session_state.current_page = "Login"
                    streamlit.rerun()
            else:
                streamlit.error("⚠️ Both fields are required!")

    # Login Page
    elif streamlit.session_state.current_page == "Login":
        streamlit.subheader("🔑 Reauthorization Required")

        if time.time() - streamlit.session_state.last_attempt_time < 10 and streamlit.session_state.failed_attempts >= 3:
            remaining_time = int(10 - (time.time() - streamlit.session_state.last_attempt_time))
            streamlit.warning(f"🕒 Please wait {remaining_time} seconds before trying again.")
        else:
            login_pass = streamlit.text_input("Enter Master Password:", type="password")

            if streamlit.button("Login"):
                if login_pass == "admin123":
                    reset_failed_attempts()
                    streamlit.success("✅ Reauthorized successfully!")
                    streamlit.session_state.current_page = "Home"
                    streamlit.rerun()
                else:
                    streamlit.error("❌ Incorrect password!")

    # Footer
    streamlit.markdown("---")
    streamlit.markdown("🔐 Secure Data Encryption System | Educational Project")

# ---------------------- Entry Point ----------------------

if __name__ == "__main__":
    main()