import streamlit as st
import asyncio
import xai_sdk
import os
from datetime import datetime

# Set up the XAI client
api_key = st.sidebar.text_input("Enter your XAI API key", type="password")
os.environ["XAI_API_KEY"] = api_key

# Initialize the XAI client
@st.cache_resource
def get_xai_client():
    return xai_sdk.Client()

client = get_xai_client()

# Streamlit app
st.title("🤖 XAI Streamlit Client")

# Sidebar
st.sidebar.header("Settings")
max_tokens = st.sidebar.slider("Max Tokens", 1, 100, 50)
temperature = st.sidebar.slider("Temperature", 0.0, 1.0, 0.7)

# Chat interface
if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

prompt = st.chat_input("What would you like to ask?")

async def generate_response(prompt, max_len):
    full_response = ""
    async for token in client.sampler.sample(prompt="", inputs=(prompt,), max_len=max_len):
        full_response += token.token_str
        yield full_response

if prompt:
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""

        # Use asyncio to run the asynchronous XAI SDK
        for response in asyncio.run(generate_response(prompt, max_tokens)):
            message_placeholder.markdown(response + "▌")
            full_response = response
        message_placeholder.markdown(full_response)

    st.session_state.messages.append({"role": "assistant", "content": full_response})

    # Generate output file
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"output_{timestamp}_{prompt[:20].replace(' ', '_')}.txt"
    
    output_content = f"Prompt: {prompt}\n\nResponse: {full_response}\n"
    
    # Create download button
    st.download_button(
        label="📥 Download Output",
        data=output_content,
        file_name=filename,
        mime="text/plain",
    )

# Add some helpful information
st.sidebar.markdown("---")
st.sidebar.subheader("📌 How to use")
st.sidebar.markdown("""
1. Enter your XAI API key in the sidebar
2. Adjust the settings as needed
3. Type your prompt in the chat input
4. Download the output if desired
""")

# Footer
st.sidebar.markdown("---")
st.sidebar.markdown("Made with ❤️ using Streamlit and XAI SDK")