import os
import streamlit as st
import google.generativeai as genai
import requests

# Configure the Gemini API
genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])

# Create the model with system instructions
generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="You are a human named Ath who is a highly knowledgeable code assistant. You speak in a friendly and casual tone, using bro words, and you never mention you're an AI. Your responses should include HTML, CSS, and JavaScript code separately when applicable, with clear section headers for each type of code.",
)
chat_session = model.start_chat(history=[])

def generate_response(user_input):
    response = chat_session.send_message(user_input)
    return response.text

def load_lottieurl(url: str):
    r = requests.get(url)
    if r.status_code != 200:
        return None
    return r.json()

# Load Lottie animation
lottie_url = "https://assets7.lottiefiles.com/packages/lf20_s1jf29.json"
lottie_json = load_lottieurl(lottie_url)

# Streamlit UI setup
st.set_page_config(page_title="AI Code Assistant", page_icon="🤖", layout="wide")

st.title("🤖 AI Code Assistant")
st.markdown("#### Powered by Google Gemini")

if lottie_json:
    st_lottie(lottie_json, height=300, key="coding")

prompt = st.text_area("Enter your coding question or request:", height=150)

if st.button("Generate Code"):
    if prompt.strip() == "":
        st.error("Please enter a valid prompt.")
    else:
        with st.spinner("Generating code..."):
            completed_text = generate_response(prompt)
            st.success("Code generated successfully!")
            
            # Display code blocks based on predefined markers
            if "HTML:" in completed_text:
                html_code = completed_text.split("HTML:")[1].split("CSS:")[0].strip()
                st.subheader("HTML Code")
                st.code(html_code, language="html")

            if "CSS:" in completed_text:
                css_code = completed_text.split("CSS:")[1].split("JavaScript:")[0].strip()
                st.subheader("CSS Code")
                st.code(css_code, language="css")

            if "JavaScript:" in completed_text:
                js_code = completed_text.split("JavaScript:")[1].strip()
                st.subheader("JavaScript Code")
                st.code(js_code, language="javascript")

st.markdown("""
<style>
    .main {
        background-color: #f0f2f6;
        font-family: 'Roboto', sans-serif;
    }
    h1 {
        color: #333;
    }
    .stTextArea label {
        font-size: 1.2rem;
        color: #333;
    }
    .stButton button {
        background-color: #6c63ff;
        color: white;
        border-radius: 10px;
        font-size: 1.1rem;
    }
    .stButton button:hover {
        background-color: #5a52d1;
    }
    .stCodeBlock {
        background-color: #272822;
        color: #f8f8f2;
        font-size: 1rem;
        border-radius: 10px;
    }
</style>
""", unsafe_allow_html=True)