# 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 Neon Theme
def inject_custom_css():
    st.markdown("""
    <style>
        .main { background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); }
        .stTextInput input { color: #4fffb0 !important; background: rgba(0,0,0,0.3) !important; }
        .stButton>button { background: rgba(79,255,176,0.2) !important; border: 2px solid #4fffb0 !important; }
        .code-editor { border: 2px solid #4fffb0; border-radius: 10px; padding: 1rem; background: rgba(0,0,0,0.5); }
    </style>
    """, unsafe_allow_html=True)

inject_custom_css()

# Initialize Groq Client
def get_groq_client():
    api_key = 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, model, client):
    try:
        completion = client.chat.completions.create(
            model=model,
            messages=[{
                "role": "user",
                "content": f"Generate Python 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 🚀")
    
    col1, col2 = st.columns([3, 1])
    with col1:
        query = st.text_input("Describe your coding requirement:")
    
    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, PRIMARY_MODEL, client)
            if not code:
                st.warning("Trying backup model...")
                code = generate_code(query, 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.markdown(f"```python\n{st.session_state.generated_code}\n```")

if __name__ == "__main__":
    main()