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

# --- Constants ---
PRIMARY_MODEL = "qwen-2.5-coder-32b"
SUPPORTED_LANGUAGES = ["Python", "JavaScript", "Java", "C++", "Go", "Rust"]
LOTTIE_URL = "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):
    response = requests.get(url)
    return response.json() if response.status_code == 200 else None

# --- Custom CSS ---
def inject_css():
    st.markdown(f"""
    <style>
        .main {{
            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
            color: #ffffff;
        }}
        
        .stTextInput input {{
            background: rgba(0,0,0,0.3) !important;
            border: 2px solid #4fffb0 !important;
            color: #4fffb0 !important;
        }}
        
        .code-editor {{
            border: 2px solid #4fffb0;
            border-radius: 15px;
            padding: 1rem;
            background: rgba(0,0,0,0.5) !important;
        }}
        
        .download-btn {{
            background: linear-gradient(45deg, #4fffb0, #2bd2ff) !important;
            color: #1a1a2e !important;
            border: none !important;
            border-radius: 25px;
            padding: 12px 24px;
        }}
    </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:
        response = client.chat.completions.create(
            model=PRIMARY_MODEL,
            messages=[{
                "role": "user",
                "content": f"Generate {lang} code for: {query}. Include comments and error handling."
            }],
            temperature=0.7,
            max_tokens=4096
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"Error: {str(e)}")
        return None

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

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

    # --- Generation Flow ---
    if st.button("🚀 Generate Code", use_container_width=True):
        if user_query:
            with st.spinner(f"🧠 Crafting {lang} code..."):
                code = generate_code(client, user_query, lang)
                if code:
                    st.session_state.code = code
                    st.balloons()
                    
                    with st.expander(f"✨ Generated Code", expanded=True):
                        st.code(code)
                        ext = {"Python": "py", "JavaScript": "js"}.get(lang, "txt")
                        st.markdown(get_download_link(code, f"code.{ext}"), unsafe_allow_html=True)
        else:
            st.warning("Please enter a requirement!")

if __name__ == "__main__":
    main()