app.py

# app.py
import os
from groq import Groq
import time
import json
import zipfile
import io
import streamlit as st
from streamlit_lottie import st_lottie
from streamlit_extras.stylable_container import stylable_container

# --- Constants ---
PRIMARY_MODEL = "qwen-2.5-coder-32b"
BACKUP_MODEL = "llama3-70b-8192"
LANGUAGES = ["Python", "JavaScript", "Java", "C++", "Go", "Rust"]
THEMES = ["Dracula", "Monokai", "Nord", "Solarized", "One Dark"]

# --- Load Assets ---
def load_lottie(filepath):
    with open(filepath, "r") as f:
        return json.load(f)

# --- Initialize Clients ---
@st.cache_resource
def init_groq():
    return Groq(api_key=os.getenv("GROQ_API_KEY"))

# --- UI Config ---
st.set_page_config(
    page_title="CodeForge AI",
    page_icon="⚡",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Custom CSS ---
def inject_css():
    st.markdown(f"""
    <style>
        [data-testid="stAppViewContainer"] {{
            background: linear-gradient(135deg, #0f0c29 0%, #302b63 50%, #24243e 100%);
        }}
        .stTextInput input {{
            color: #00ffaa !important;
            background: rgba(0,0,0,0.3) !important;
            border: 2px solid #00ffaa !important;
            box-shadow: 0 0 10px rgba(0,255,170,0.5);
        }}
        .stButton>button {{
            background: linear-gradient(90deg, #00ffaa 0%, #0088ff 100%) !important;
            color: #111 !important;
            font-weight: bold !important;
            border: none !important;
            box-shadow: 0 0 15px rgba(0,255,170,0.7);
            transition: all 0.3s ease !important;
        }}
        .stButton>button:hover {{
            transform: scale(1.05);
            box-shadow: 0 0 20px rgba(0,255,170,0.9);
        }}
        .code-header {{
            color: #00ffaa;
            font-family: 'Fira Code', monospace;
            text-shadow: 0 0 10px rgba(0,255,170,0.7);
        }}
        [data-testid="stSidebar"] {{
            background: rgba(16, 18, 33, 0.8) !important;
            backdrop-filter: blur(10px);
            border-right: 1px solid rgba(0,255,170,0.3);
        }}
    </style>
    """, unsafe_allow_html=True)

inject_css()

# --- Animation ---
def render_hero():
    col1, col2 = st.columns([2, 1])
    with col1:
        st.markdown("""
        # <span class="code-header">⚡ CodeForge AI</span>
        ### _Generate • Optimize • Deploy_
        """, unsafe_allow_html=True)
    with col2:
        st_lottie(load_lottie("https://assets9.lottiefiles.com/packages/lf20_2znxgjyt.json"), 
                 height=150, key="hero")

# --- Main Generator ---
def generate_code(client, prompt, language, model):
    messages = [{
        "role": "system",
        "content": f"You are an expert {language} developer. Generate clean, efficient, production-ready code."
    }, {
        "role": "user",
        "content": prompt
    }]
    
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.7,
        max_tokens=4000,
        top_p=0.95,
        stream=True
    )
    
    return response

# --- UI Components ---
def language_selector():
    with st.sidebar:
        with stylable_container(
            key="lang-select",
            css_styles="""
                {
                    border: 2px solid #00ffaa;
                    border-radius: 10px;
                    padding: 10px;
                    background: rgba(0,0,0,0.2);
                }
            """
        ):
            return st.selectbox("🗝️ Select Language", LANGUAGES, index=0)

def theme_selector():
    with st.sidebar:
        with stylable_container(
            key="theme-select",
            css_styles="""
                {
                    border: 2px solid #0088ff;
                    border-radius: 10px;
                    padding: 10px;
                    background: rgba(0,0,0,0.2);
                }
            """
        ):
            return st.selectbox("🎨 Editor Theme", THEMES, index=0)

# --- Download Handler ---
def create_download_zip(code, language):
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "a") as zf:
        zf.writestr(f"generated_code.{language.lower()}", code)
    return zip_buffer.getvalue()

# --- Main App ---
def main():
    client = init_groq()
    render_hero()
    
    # Sidebar Controls
    language = language_selector()
    theme = theme_selector()
    
    # Main Input
    with st.form(key="input_form"):
        prompt = st.text_area("💡 Describe your coding task:", height=150,
                            placeholder="e.g. 'Create a REST API with FastAPI and MongoDB'")
        
        col1, col2, col3 = st.columns([1,1,2])
        with col1:
            generate_btn = st.form_submit_button("✨ Generate Code")
        with col2:
            explain_btn = st.form_submit_button("📚 Explain")
        with col3:
            optimize_btn = st.form_submit_button("⚡ Optimize")
    
    # Generation Logic
    if generate_btn and prompt:
        with st.spinner("🚀 Generating your code..."):
            try:
                response = generate_code(client, prompt, language, PRIMARY_MODEL)
                
                code_container = st.empty()
                full_response = []
                
                for chunk in response:
                    if chunk.choices[0].delta.content:
                        full_response.append(chunk.choices[0].delta.content)
                        current_code = "".join(full_response)
                        code_container.code(current_code, language=language.lower(), line_numbers=True)
                
                st.session_state.generated_code = "".join(full_response)
                
                # Download Button
                zip_data = create_download_zip(st.session_state.generated_code, language)
                st.download_button(
                    label="📥 Download Code",
                    data=zip_data,
                    file_name=f"codeforge_{int(time.time())}.zip",
                    mime="application/zip"
                )
                
            except Exception as e:
                st.error(f"🚨 Generation failed: {str(e)}")
                st.info("Attempting with backup model...")
                try:
                    response = generate_code(client, prompt, language, BACKUP_MODEL)
                    # ... [backup handling logic]
                except Exception as e:
                    st.error("All models failed. Please try again later.")

if __name__ == "__main__":
    main()