app.py

import os
import streamlit as st
from groq import Groq

# Configuration
PRIMARY_MODEL = "qwen-2.5-coder-32b"
BACKUP_MODEL = "llama3-70b-8192"

# Streamlit UI Config
st.set_page_config(
    page_title="AI Code Generator",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Enhanced UI
def inject_custom_css():
    st.markdown("""
    <style>
        .main { background: linear-gradient(135deg, #0f0f0f 0%, #1a1a2e 100%); color: #fff; }
        .stTextInput input { color: #4fffb0 !important; background: rgba(0,0,0,0.3) !important; }
        .stButton>button { background: #4fffb0 !important; border: 2px solid #00ffcc !important; color: black; }
        .code-box { border: 2px solid #4fffb0; border-radius: 10px; padding: 1rem; background: rgba(0,0,0,0.5); color: #fff; }
        .stSelectbox div { background: rgba(0,0,0,0.3) !important; color: #4fffb0 !important; }
    </style>
    """, unsafe_allow_html=True)

inject_custom_css()

# Initialize Groq Client
def get_groq_client():
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        st.error("GROQ_API_KEY not found! Please set it in environment variables or secrets.")
        return None
    return Groq(api_key=api_key)

# Code Generation Functions
def generate_code(query, language, model, client):
    try:
        completion = client.chat.completions.create(
            model=model,
            messages=[{
                "role": "user",
                "content": f"Generate {language} code for: {query}. Include comments and error handling."
            }],
            temperature=0.6,
            max_tokens=4096,
            top_p=0.95
        )
        return completion.choices[0].message.content
    except Exception as e:
        st.error(f"Error generating code: {str(e)}")
        return None

# Main App Interface
def main():
    client = get_groq_client()
    if not client:
        return

    st.title("🔥 AI Code Generator 🚀")
    st.subheader("Generate, Edit & Save Your Code Instantly!")
    
    col1, col2 = st.columns([3, 1])
    
    with col1:
        query = st.text_input("Describe your coding requirement:")
        language = st.selectbox("Select Programming Language:", ["Python", "JavaScript", "Java", "C++", "Go"])
    
    if 'generated_code' not in st.session_state:
        st.session_state.generated_code = None
    
    if col2.button("🔄 New Code"):
        st.session_state.generated_code = None
    
    if query and not st.session_state.generated_code:
        with st.spinner("Generating code..."):
            code = generate_code(query, language, PRIMARY_MODEL, client)
            if not code:
                st.warning("Trying backup model...")
                code = generate_code(query, language, BACKUP_MODEL, client)
            
            if code:
                st.session_state.generated_code = code
            else:
                st.error("Failed to generate code. Please try again.")
    
    if st.session_state.generated_code:
        with st.expander("📜 Generated Code", expanded=True):
            st.code(st.session_state.generated_code, language.lower())
            
            col3, col4 = st.columns([1, 1])
            with col3:
                st.download_button("💾 Download Code", st.session_state.generated_code, file_name=f"generated.{language.lower()}")
            with col4:
                st.button("📋 Copy Code", on_click=lambda: st.session_state.update({'clipboard': st.session_state.generated_code}))

if __name__ == "__main__":
    main()