import os
import streamlit as st
from groq import Groq
from dotenv import load_dotenv

# Load API key
load_dotenv()
API_KEY = os.getenv("GROQ_API_KEY")

# Initialize Groq client
client = Groq(api_key=API_KEY)

# Streamlit App Title
st.title("🧠 Brain-to-Text AI Keyboard")

# User input: Simulated brain signal (text description)
st.subheader("Think of a sentence, then describe it in words:")
thought_input = st.text_area("Describe your thought:", height=100)

# Virtual keyboard (optional)
st.subheader("Or use the virtual keyboard to refine your thought:")
if "typed_text" not in st.session_state:
    st.session_state.typed_text = ""

# Define keyboard layout
keyboard_layout = [
    ["Q", "W", "E", "R", "T", "Y", "U", "I", "O", "P"],
    ["A", "S", "D", "F", "G", "H", "J", "K", "L"],
    ["Z", "X", "C", "V", "B", "N", "M", "←"],
    ["Space"]
]

# Display keyboard
for row in keyboard_layout:
    cols = st.columns(len(row))
    for i, key in enumerate(row):
        if cols[i].button(key):
            if key == "←":
                st.session_state.typed_text = st.session_state.typed_text[:-1]
            elif key == "Space":
                st.session_state.typed_text += " "
            else:
                st.session_state.typed_text += key

# Display user-typed text
st.text_area("Refined Thought:", value=st.session_state.typed_text, height=100)

# Generate AI-based Brain-to-Text Conversion
if st.button("Convert Thought to Text"):
    full_input = thought_input.strip() + " " + st.session_state.typed_text.strip()
    if full_input.strip():
        try:
            response = client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[{"role": "user", "content": full_input}]
            )
            st.write("### AI Interpreted Thought:")
            st.write(response.choices[0].message.content)
        except Exception as e:
            st.error(f"Error: {e}")
    else:
        st.warning("Please describe a thought before converting.")