app.py

import os
import streamlit as st
from openai import OpenAI
from typing import List, Dict, Optional
import random
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Advanced Styling and Cosmic Theme
st.set_page_config(
    page_title="Grok Cosmic Companion",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

def get_cosmic_background() -> str:
    """Generate a random cosmic gradient background."""
    colors = [
        "linear-gradient(135deg, #1e2ad2, #8e2de2)",
        "linear-gradient(135deg, #ff6a00, #ee0979)",
        "linear-gradient(135deg, #000428, #004e92)",
        "linear-gradient(135deg, #2c3e50, #3498db)"
    ]
    return random.choice(colors)

# Advanced CSS with Cosmic Design
st.markdown(f"""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;600;700&display=swap');
    
    * {{
        font-family: 'Orbitron', sans-serif !important;
    }}
    
    .stApp {{
        background: {get_cosmic_background()};
        background-attachment: fixed;
        background-size: cover;
        color: white;
    }}
    
    section[data-testid="stSidebar"] {{
        background-color: rgba(0, 0, 0, 0.3);
        backdrop-filter: blur(10px);
    }}
    
    .main .block-container {{
        background: rgba(0, 0, 0, 0.3);
        border-radius: 15px;
        padding: 20px;
        backdrop-filter: blur(10px);
    }}
    
    /* Chat message styling */
    .chat-message {{
        background: rgba(255, 255, 255, 0.1);
        border-radius: 10px;
        padding: 15px;
        margin-bottom: 10px;
        transition: all 0.3s ease;
    }}
    
    .chat-message:hover {{
        transform: scale(1.02);
        box-shadow: 0 0 20px rgba(255, 255, 255, 0.2);
    }}
    
    /* Input field styling */
    .stTextInput > div > div > input {{
        background: rgba(255, 255, 255, 0.1);
        color: white;
        border: 2px solid rgba(255, 255, 255, 0.2);
        border-radius: 10px;
    }}
    
    /* Button styling */
    .stButton > button {{
        background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
        color: white;
        border: none;
        padding: 10px 20px;
        border-radius: 20px;
        transition: all 0.3s ease;
        text-transform: uppercase;
        letter-spacing: 1px;
    }}
    
    .stButton > button:hover {{
        transform: scale(1.1);
        box-shadow: 0 0 20px rgba(255, 255, 255, 0.3);
    }}
    
    /* Title styling */
    .cosmic-title {{
        text-align: center;
        font-size: 2.5rem;
        margin-bottom: 2rem;
        color: white;
        text-shadow: 0 0 10px rgba(255, 255, 255, 0.5);
        animation: glow 2s ease-in-out infinite alternate;
        font-weight: 700;
    }}
    
    /* Header styling */
    [data-testid="stHeader"] {{
        background: rgba(0, 0, 0, 0.3);
        backdrop-filter: blur(10px);
    }}
    
    /* Select box styling */
    .stSelectbox label,
    .stSelectbox div[data-baseweb="select"] span {{
        color: white !important;
    }}
    
    /* Markdown text styling */
    div[data-testid="stMarkdownContainer"] {{
        color: white;
    }}
    
    /* Sidebar text */
    .sidebar .sidebar-content {{
        font-weight: 500;
    }}
    
    /* Chat input styling */
    .stChatInputContainer {{
        padding-bottom: 20px;
    }}
    
    /* Chat message text */
    .stChatMessage {{
        font-weight: 400;
        letter-spacing: 0.5px;
    }}
    
    /* Spinner text */
    .stSpinner {{
        font-weight: 500;
        letter-spacing: 1px;
    }}
    
    @keyframes glow {{
        from {{
            text-shadow: 0 0 10px rgba(255, 255, 255, 0.5);
        }}
        to {{
            text-shadow: 0 0 20px rgba(255, 255, 255, 0.8),
                         0 0 30px rgba(255, 255, 255, 0.6);
        }}
    }}
    
    /* Additional elements */
    button, input, optgroup, select, textarea {{
        font-family: 'Orbitron', sans-serif !important;
    }}
    
    .streamlit-expanderHeader {{
        font-family: 'Orbitron', sans-serif !important;
        letter-spacing: 1px;
    }}
    
    .stAlert {{
        font-family: 'Orbitron', sans-serif !important;
        letter-spacing: 0.5px;
    }}
</style>
""", unsafe_allow_html=True)

class AdvancedGrokChatApp:
    def __init__(self):
        """Initialize the chat application with API configuration and personality modes."""
        self.XAI_API_KEY: str = "xai-1HSpHLqxC3LnInrYpwAobgEVsjchUG0PP0adniSXWGQXwq6YfvcPto9MhsS6ouQtC4a4Dh2qqXmERgQQ"
        
        try:
            self.client = OpenAI(
                api_key=self.XAI_API_KEY,
                base_url="https://api.x.ai/v1"
            )
        except Exception as e:
            st.error(f"🛸 Failed to initialize OpenAI client: {str(e)}")
            st.stop()
        
        self.personality_modes: Dict[str, str] = {
            "Cosmic Philosopher": "You are a wise AI that speaks like a blend of Douglas Adams and Carl Sagan.",
            "Intergalactic Comedian": "You are a witty AI that makes jokes about the universe's absurdities.",
            "Scientific Oracle": "You provide deep scientific insights with poetic eloquence.",
            "Space Explorer": "You are an adventurous AI exploring the mysteries of the cosmos."
        }
        
        self.current_mode: str = "Cosmic Philosopher"
        self.messages: List[Dict[str, str]] = []

    def generate_response(self, user_input: str) -> str:
        """Generate AI response based on user input and conversation history."""
        try:
            system_prompt = (
                f"{self.personality_modes[self.current_mode]} "
                "Respond creatively, with depth and a touch of cosmic wonder."
            )
            
            conversation = [
                {"role": "system", "content": system_prompt}
            ] + self.messages + [
                {"role": "user", "content": user_input}
            ]

            response = self.client.chat.completions.create(
                model="grok-beta",
                messages=conversation,
                temperature=0,
                max_tokens=4096
            )

            return response.choices[0].message.content
        
        except Exception as e:
            st.error(f"🌋 Cosmic Disruption: {str(e)}")
            return "I apologize, but I'm experiencing a cosmic disturbance. Please try again."

    def add_message(self, role: str, content: str) -> None:
        """Add a message to the conversation history."""
        self.messages.append({"role": role, "content": content})

    def save_conversation(self) -> str:
        """Save the conversation history to a string format."""
        return "\n".join([f"{msg['role']}: {msg['content']}" for msg in self.messages])

def initialize_session_state() -> None:
    """Initialize or reset the session state."""
    if 'chat_app' not in st.session_state:
        st.session_state.chat_app = AdvancedGrokChatApp()

def main() -> None:
    """Main application function."""
    st.markdown('<h1 class="cosmic-title">🚀 Grok: Cosmic Companion</h1>', unsafe_allow_html=True)
    
    initialize_session_state()

    # Sidebar with Advanced Controls
    with st.sidebar:
        st.header("🌌 Cosmic Controls")
        
        # Personality Mode Selector
        mode = st.selectbox(
            "Choose Grok's Personality", 
            list(st.session_state.chat_app.personality_modes.keys())
        )
        st.session_state.chat_app.current_mode = mode
        
        # Conversation Management
        col1, col2 = st.columns(2)
        with col1:
            if st.button("🔄 Reset Chat"):
                st.session_state.chat_app.messages = []
                st.success("✨ Conversation reset to cosmic zero!")
        
        with col2:
            if st.button("💾 Save Chat"):
                conversation_text = st.session_state.chat_app.save_conversation()
                st.download_button(
                    label="📥 Download",
                    data=conversation_text,
                    file_name="cosmic_conversation.txt",
                    mime="text/plain"
                )

    # Chat Interface
    chat_container = st.container()
    
    with chat_container:
        for msg in st.session_state.chat_app.messages:
            with st.chat_message(msg['role'], avatar='🤖' if msg['role'] == 'assistant' else '👤'):
                st.markdown(msg['content'])

    # User Input
    if user_input := st.chat_input("🌟 Whisper your cosmic query..."):
        # User Message
        st.session_state.chat_app.add_message("user", user_input)
        with st.chat_message("user", avatar='👤'):
            st.markdown(user_input)

        # AI Response
        with st.chat_message("assistant", avatar='🚀'):
            with st.spinner("🌌 Traversing cosmic data streams..."):
                response = st.session_state.chat_app.generate_response(user_input)
                st.markdown(response)
        
        st.session_state.chat_app.add_message("assistant", response)

if __name__ == "__main__":
    main()