import streamlit as st
from groq import Groq
from langgraph.graph import StateGraph, END
from langchain.prompts import PromptTemplate
import os
from typing import Dict, TypedDict
from datetime import datetime

# Initialize Groq client (API key from environment or Streamlit secrets)
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY", st.secrets.get("GROQ_API_KEY", "")))

# Available open-source models from Groq (as of April 2025, based on current offerings)
OPEN_SOURCE_MODELS = {
    "LLaMA3 8B": "llama3-8b-8192",
    "LLaMA3 70B": "llama3-70b-8192",
    "Mixtral 8x7B": "mixtral-8x7b-32768",
    "Gemma 7B": "gemma-7b-it"
}

# Define the state for LangGraph
class ChatState(TypedDict):
    messages: list
    response: str
    selected_model: str

# Define the chatbot node for LangGraph
def chatbot_node(state: ChatState) -> ChatState:
    user_input = state["messages"][-1]["content"]
    selected_model = state["selected_model"]
    prompt = PromptTemplate(
        input_variables=["input"],
        template="You are a helpful assistant. Respond to this: {input}"
    ).format(input=user_input)
    
    # Call Groq API with the selected model
    response = groq_client.chat.completions.create(
        model=selected_model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=500,
        temperature=0.7
    )
    
    state["response"] = response.choices[0].message.content
    return state

# Build the LangGraph workflow
workflow = StateGraph(ChatState)
workflow.add_node("chatbot", chatbot_node)
workflow.set_entry_point("chatbot")
workflow.add_edge("chatbot", END)
app = workflow.compile()

# Streamlit UI
def main():
    # Set page config for a wider layout and custom title
    st.set_page_config(page_title="Open-Source Chatbot", layout="wide", initial_sidebar_state="expanded")

    # Custom CSS for a classy UI
    st.markdown("""
        <style>
        .main {background-color: #f9f9f9; padding: 20px; border-radius: 10px;}
        .stTextInput > div > div > input {border-radius: 8px; padding: 10px; font-size: 16px;}
        .chat-message {padding: 15px; margin: 10px 0; border-radius: 10px; box-shadow: 0 2px 5px rgba(0,0,0,0.1);}
        .user-message {background-color: #e3f2fd; color: #0d47a1;}
        .bot-message {background-color: #e8f5e9; color: #1b5e20;}
        .sidebar .sidebar-content {background-color: #ffffff; padding: 20px; border-radius: 10px;}
        h1 {font-family: 'Arial', sans-serif; color: #1a237e;}
        .stRadio > label {font-size: 16px; padding: 8px; background-color: #f0f0f0; border-radius: 5px; margin: 5px 0;}
        .stRadio > label:hover {background-color: #e0e0e0;}
        </style>
    """, unsafe_allow_html=True)

    # Sidebar for settings and model selection
    with st.sidebar:
        st.header("Settings")
        st.info("Select a model and enter your Groq API key to start.")
        
        # Model selection with radio buttons
        selected_model_name = st.radio(
            "Choose a Model",
            options=list(OPEN_SOURCE_MODELS.keys()),
            index=0,
            help="Pick an open-source model to power the chatbot."
        )
        selected_model = OPEN_SOURCE_MODELS[selected_model_name]

        # API key input
        api_key_input = st.text_input("Groq API Key", type="password", value=os.getenv("GROQ_API_KEY", ""))
        if api_key_input:
            os.environ["GROQ_API_KEY"] = api_key_input
            st.success("API Key set successfully!")
        
        st.markdown("---")
        st.write(f"Current Time: {datetime.now().strftime('%H:%M:%S %d-%m-%Y')}")
        st.write(f"Selected Model: **{selected_model_name}**")

    # Main chat interface
    st.title("Open-Source Chatbot")
    st.write("A sleek chatbot powered by open-source models via Groq.")

    # Initialize session state for chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Display chat history
    for msg in st.session_state.messages:
        if msg["role"] == "user":
            st.markdown(f'<div class="chat-message user-message">You: {msg["content"]}</div>', unsafe_allow_html=True)
        else:
            st.markdown(f'<div class="chat-message bot-message">Bot: {msg["content"]}</div>', unsafe_allow_html=True)

    # User input
    with st.form(key="chat_form", clear_on_submit=True):
        user_input = st.text_input("Type your message here...", key="input")
        submit_button = st.form_submit_button(label="Send")

    # Process input and get response
    if submit_button and user_input:
        if not os.getenv("GROQ_API_KEY"):
            st.error("Please provide a Groq API key in the sidebar!")
        else:
            st.session_state.messages.append({"role": "user", "content": user_input})
            state = {
                "messages": st.session_state.messages,
                "response": "",
                "selected_model": selected_model
            }
            result = app.invoke(state)
            st.session_state.messages.append({"role": "assistant", "content": result["response"]})
            st.rerun()  # Rerun to update the UI

if __name__ == "__main__":
    main()