import streamlit as st
import streamlit_authenticator as stauth
import yaml
from yaml.loader import SafeLoader
from groq import Groq
import base64

# ---------------- AUTHENTICATION SETUP ----------------
# Load credentials from a YAML file
CREDENTIALS = {
    "usernames": {
        "user1": {"name": "John Doe", "password": stauth.Hasher(["password123"]).generate()[0]},
        "user2": {"name": "Jane Doe", "password": stauth.Hasher(["securepass"]).generate()[0]},
    }
}

# Convert credentials to YAML format
with open("credentials.yaml", "w") as file:
    yaml.dump({"credentials": CREDENTIALS}, file)

# Load credentials
with open("credentials.yaml") as file:
    config = yaml.load(file, Loader=SafeLoader)

# Initialize authenticator
authenticator = stauth.Authenticate(
    config["credentials"], "finance_chat", "abcdef", cookie_expiry_days=30
)

# Show login widget
name, authentication_status, username = authenticator.login("Login", "main")

# ---------------- USER AUTHENTICATION HANDLING ----------------
if authentication_status:
    authenticator.logout("Logout", "sidebar")
    st.sidebar.write(f"Welcome, **{name}**! 🎉")

    # ---------------- CHATBOT SETUP ----------------
    # Set up Groq API client
    client = Groq(api_key="your-groq-api-key")

    st.set_page_config(page_title="💰 Finance & Banking Chatbot", layout="wide")
    st.title("💳 Finance & Banking Chatbot 🤵")

    # System prompt to enforce finance-related responses
    SYSTEM_PROMPT = (
        "You are an expert financial assistant. Your role is to answer ONLY finance-related topics, "
        "including banking, investments, loans, credit cards, budgeting, and economic trends. "
        "If a user asks about unrelated topics, respond: "
        "'I'm here to assist with financial topics only. 💰'"
    )

    # Initialize session state for chat history
    if "messages" not in st.session_state:
        st.session_state.messages = [
            {"role": "assistant", "content": "Hello! How can I assist with your finance questions? 💳"}
        ]

    # Display previous messages
    for msg in st.session_state.messages:
        avatar = "👦🏻" if msg["role"] == "user" else "🤵"
        st.chat_message(msg["role"], avatar=avatar).write(msg["content"])

    # User input
    user_input = st.chat_input("Ask me about finance, banking, investments, etc. 📈")

    if user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})
        st.chat_message("user", avatar="👦🏻").write(user_input)

        # Get response from Groq API
        response = client.chat.completions.create(
            messages=[{"role": "system", "content": SYSTEM_PROMPT}] + st.session_state.messages,
            model="llama-3.3-70b-versatile",
            max_tokens=200,
        )

        bot_reply = response.choices[0].message.content
        st.session_state.messages.append({"role": "assistant", "content": bot_reply})
        st.chat_message("assistant", avatar="🤵").write(bot_reply)

    # ---------------- SHARING OPTION ----------------
    st.sidebar.title("🔗 Share Your Conversation")
    
    # Convert chat history to text
    chat_text = "\n".join(
        [f"{'User' if msg['role'] == 'user' else 'Bot'}: {msg['content']}" for msg in st.session_state.messages]
    )

    # Encode chat history to a downloadable text file
    b64_chat = base64.b64encode(chat_text.encode()).decode()
    href = f'<a href="data:file/txt;base64,{b64_chat}" download="chat_history.txt">📥 Download Chat History</a>'
    st.sidebar.markdown(href, unsafe_allow_html=True)

else:
    st.warning("Please log in to access the chatbot.")