# app.py
import os
import time
import streamlit as st
from groq import Groq
from streamlit_lottie import st_lottie
import json
import base64

# --- Constants ---
PRIMARY_MODEL = "qwen-2.5-coder-32b"
SUPPORTED_LANGUAGES = ["Python", "JavaScript", "Java", "C++", "Go", "Rust"]
LOTTIE_ANIMATION = "https://assets3.lottiefiles.com/packages/lf20_hi95bvmx/ai.json"

# --- Setup ---
st.set_page_config(
    page_title="CodeForge AI",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Lottie Animation Loader ---
def load_lottie(url):
    with open("animation.json") as f:
        return json.load(f)

# --- Custom CSS with Animations ---
def inject_css():
    st.markdown(f"""
    <style>
        @keyframes glow {{
            0% {{ box-shadow: 0 0 5px #4fffb0; }}
            50% {{ box-shadow: 0 0 20px #4fffb0; }}
            100% {{ box-shadow: 0 0 5px #4fffb0; }}
        }}
        
        .main {{
            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
            color: #ffffff;
        }}
        
        .stTextInput input {{
            animation: glow 2s infinite;
            background: rgba(0,0,0,0.3) !important;
            border: 2px solid #4fffb0 !important;
        }}
        
        .code-editor {{
            border: 2px solid #4fffb0;
            border-radius: 15px;
            padding: 1rem;
            background: rgba(0,0,0,0.5) !important;
            animation: glow 2s infinite;
        }}
        
        .download-btn {{
            background: linear-gradient(45deg, #4fffb0, #2bd2ff) !important;
            color: #1a1a2e !important;
            border: none !important;
            border-radius: 25px;
            padding: 12px 24px;
            font-weight: bold;
            transition: transform 0.3s !important;
        }}
        
        .download-btn:hover {{
            transform: scale(1.05);
        }}
    </style>
    """, unsafe_allow_html=True)

# --- File Download Generator ---
def get_download_link(code, filename):
    b64 = base64.b64encode(code.encode()).decode()
    return f'<a href="data:file/txt;base64,{b64}" download="{filename}" class="download-btn">⬇️ Download {filename}</a>'

# --- Model Client ---
@st.cache_resource
def init_groq():
    return Groq(api_key=os.environ.get("GROQ_API_KEY"))

# --- Code Generation Core ---
def generate_code(client, query, lang):
    try:
        system_prompt = f"""
        You are a {lang} expert. Generate:
        1. Clean, production-ready code
        2. With error handling
        3. Proper documentation
        4. Optimized for performance
        """
        
        response = client.chat.completions.create(
            model=PRIMARY_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": query}
            ],
            temperature=0.7,
            max_tokens=4096,
            top_p=0.95
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"🚨 Generation Error: {str(e)}")
        return None

# --- UI Components ---
def main():
    inject_css()
    client = init_groq()
    
    # --- Hero Section ---
    col1, col2 = st.columns([1, 2])
    with col1:
        st_lottie(load_lottie(LOTTIE_ANIMATION), height=200)
    with col2:
        st.markdown("# CodeForge AI")
        st.markdown("### _Transform Ideas into Production Code_ ⚡")

    # --- Language Selector ---
    lang = st.sidebar.selectbox("💻 Programming Language", SUPPORTED_LANGUAGES, index=0)
    user_query = st.text_input("🎯 Describe your coding requirement:", placeholder="e.g., 'Create a REST API with JWT authentication'")

    # --- Main Generation Flow ---
    if st.button("🚀 Generate Code", use_container_width=True):
        if not user_query:
            st.warning("⚠️ Please enter a coding requirement!")
            return
            
        with st.spinner(f"🧠 Crafting {lang} magic..."):
            start_time = time.time()
            code = generate_code(client, user_query, lang)
            gen_time = time.time() - start_time
            
            if code:
                st.session_state.code = code
                st.balloons()
                
                # --- Code Display ---
                with st.expander(f"✨ Generated {lang} Code", expanded=True):
                    st.code(code, language="python" if lang == "Python" else "javascript")
                    
                    # --- Download Button ---
                    ext = {"Python": "py", "JavaScript": "js", "Java": "java", 
                          "C++": "cpp", "Go": "go", "Rust": "rs"}[lang]
                    st.markdown(get_download_link(code, f"generated_code.{ext}"), 
                               unsafe_allow_html=True)
                
                # --- Performance Metrics ---
                st.sidebar.markdown(f"""
                ### 📊 Generation Metrics
                ⏱️ **Time**: {gen_time:.2f}s  
                📝 **Tokens**: {len(code.split())}  
                🧠 **Model**: {PRIMARY_MODEL}
                """)
                
                # --- Explanation Section ---
                if st.button("📖 Explain Implementation"):
                    with st.spinner("🔍 Analyzing code structure..."):
                        explanation = client.chat.completions.create(
                            model=PRIMARY_MODEL,
                            messages=[{
                                "role": "user",
                                "content": f"Explain this {lang} code:\n\n{code}"
                            }]
                        )
                        st.markdown(f"## Code Breakdown\n{explanation.choices[0].message.content}")

if __name__ == "__main__":
    main()