import streamlit as st
import os
from groq import Groq

# UI Setup
st.set_page_config(page_title="CodeGen Pro", page_icon="💡", layout="wide")
st.markdown("<h1 style='text-align:center;'>⚡ AI Code Generator with Qwen 32B</h1>", unsafe_allow_html=True)
st.markdown("Generate clean, documented code with the power of Qwen-32B model by Groq 🚀")

# Sidebar Configuration
with st.sidebar:
    st.header("🛠️ Code Settings")
    language = st.selectbox("Programming Language", ["Python", "JavaScript", "Go", "Java", "C++", "TypeScript"])
    style = st.selectbox("Coding Style", ["Default", "Functional", "OOP", "Concise", "Verbose"])
    tags = st.text_input("🔖 Tags (comma-separated)", placeholder="API, login, JWT, MongoDB")
    temperature = st.slider("🎛 Temperature", 0.0, 1.0, 0.6)
    max_tokens = st.slider("🧠 Max Tokens", 256, 4096, 2048)
    st.markdown("✅ Make sure your `GROQ_API_KEY` is set as environment variable.")

# Prompt builder
def build_prompt(user_input, language, style, tags):
    base_instruction = f"""
You are a senior software engineer and your task is to generate clean, well-documented, and executable {language} code. Follow these requirements:

Language: {language}
Style: {style if style != "Default" else "any"}
Tags: {tags or "N/A"}

💡 Instructions:
- Use best practices in {language}
- Include comments where needed
- Avoid unnecessary explanations
- The code must be copy-paste ready

📝 User Requirement:
\"\"\"
{user_input.strip()}
\"\"\"
"""
    return base_instruction.strip()

# Streaming generator
def stream_code(prompt):
    client = Groq(api_key=os.getenv("GROQ_API_KEY"))
    response = client.chat.completions.create(
        model="qwen/qwen3-32b",
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
        max_completion_tokens=max_tokens,
        top_p=0.95,
        reasoning_effort="default",
        stream=True,
        stop=None,
    )
    full_code = ""
    for chunk in response:
        token = chunk.choices[0].delta.content or ""
        full_code += token
        yield token
    return full_code

# Main input area
user_input = st.text_area("🧠 Describe what you want the code to do:", height=200, placeholder="e.g. Create a FastAPI JWT login endpoint connected to MongoDB...")

# Action
if st.button("🚀 Generate Code"):
    if not user_input.strip():
        st.warning("Please enter a valid code request!")
    else:
        prompt = build_prompt(user_input, language, style, tags)
        st.subheader("📜 Prompt Sent to Model")
        with st.expander("🔍 View Full Prompt", expanded=False):
            st.code(prompt, language="markdown")

        st.subheader("🧾 Generated Code")
        code_area = st.empty()
        generated_text = ""
        for token in stream_code(prompt):
            generated_text += token
            code_area.code(generated_text, language=language.lower())