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

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

# Init client
client = Groq(api_key=GROQ_API_KEY)

# App UI Setup
st.set_page_config(page_title="🧠 Code Snippet Generator", page_icon="💡", layout="wide")

# Custom CSS for style
st.markdown("""
    <style>
    .main {
        background-color: #1e1e2f;
        color: #f2f2f2;
        font-family: 'Segoe UI', sans-serif;
    }
    .stTextArea textarea {
        background-color: #2e2e3f;
        color: white;
    }
    .stButton button {
        background-color: #4CAF50;
        color: white;
    }
    </style>
""", unsafe_allow_html=True)

# Sidebar
st.sidebar.image("https://cdn-icons-png.flaticon.com/512/911/911408.png", width=100)
st.sidebar.title("Quick Code Gen 💻")
st.sidebar.markdown("Turn natural language into code + dry run explanation!")

# Main interface
st.title("⚡ Instant Code Snippet Generator")
st.subheader("Just describe what you want...")

prompt = st.text_area("💬 Describe your code idea:", height=150, placeholder="e.g., Write a Python function to reverse a string")

if st.button("🚀 Generate Code"):
    if prompt.strip() == "":
        st.warning("Please enter something first.")
    else:
        full_prompt = f"""
You are an expert programmer. Based on the input below, do two things:
1. Generate the full code snippet in Python
2. Then provide a step-by-step dry run explaining how the code works.

Input: {prompt}
"""

        with st.spinner("Generating your code..."):
            response = client.chat.completions.create(
                messages=[
                    {"role": "user", "content": full_prompt}
                ],
                model="llama-3.3-70b-versatile",
            )
            output = response.choices[0].message.content
            st.markdown("### 🧩 Generated Code + Dry Run:")
            st.code(output, language="python")