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 with API key from secrets
groq_client = Groq(api_key=st.secrets.get("GROQ_API_KEY", os.getenv("GROQ_API_KEY", "")))

# Define the model
MODEL = "llama3-70b-8192"  # LLaMA3 70B via Groq

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

# Define the chatbot node for LangGraph
def chatbot_node(state: ChatState) -> ChatState:
    user_input = state["messages"][-1]["content"]
    prompt = PromptTemplate(
        input_variables=["input"],
        template="You are a helpful assistant. Respond to this: {input}"
    ).format(input=user_input)
    
    response = groq_client.chat.completions.create(
        model=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():
    st.set_page_config(page_title="LLaMA3 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;}
        </style>
    """, unsafe_allow_html=True)

    # Sidebar for settings (no API key input)
    with st.sidebar:
        st.header("Settings")
        st.info("Powered by Groq’s LLaMA3 70B model.")
        
        st.markdown("---")
        st.write(f"Current Time: {datetime.now().strftime('%H:%M:%S %d-%m-%Y')}")
        st.write("Model: **LLaMA3 70B (Groq)**")

    # Main chat interface
    st.title("LLaMA3 Chatbot")
    st.write("A sleek chatbot powered by LLaMA3 70B 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 st.secrets.get("GROQ_API_KEY") and not os.getenv("GROQ_API_KEY"):
            st.error("Groq API key not found in secrets or environment variables!")
        else:
            st.session_state.messages.append({"role": "user", "content": user_input})
            state = {
                "messages": st.session_state.messages,
                "response": ""
            }
            result = app.invoke(state)
            st.session_state.messages.append({"role": "assistant", "content": result["response"]})
            st.rerun()

if __name__ == "__main__":
    main()