import streamlit as st
from langchain_groq import ChatGroq
from groq import Groq
import os
from streamlit_ace import st_ace
import re

# Initialize the Groq client
api_key=os.getenv("GROQ_KEY")
llm = Groq(api_key=api_key)

MODELS = [
    "llama-3.1-70b-versatile",
    "llama-3.1-8b-instant",
    "llama3-groq-70b-8192-tool-use-preview",
    "llama3-groq-8b-8192-tool-use-preview",
    "llama-guard-3-8b",
    "llama3-70b-8192",
    "llama3-8b-8192",
    "mixtral-8x7b-32768",
    "gemma-7b-it",
    "gemma2-9b-it"
]

STYLES = [
    "Modern", "Vintage", "Minimalist", "Futuristic", "Retro",
    "Glassmorphism", "Neumorphism", "Material Design", "Flat Design", "Cyberpunk",
]

CSS_FRAMEWORKS = ["Tailwind CSS", "SCSS", "Vanilla CSS"]

def generate_component(user_requirement: str, model_type: str, style: str, css_framework: str):
    prompt = f"""
    As a world-class Frontend Engineer, create a cutting-edge UI component based on:
    
    User Requirement: {user_requirement}
    Desired Style: {style}
    CSS Framework: {css_framework}
    
    Provide separate HTML and CSS/SCSS/Tailwind classes as applicable.
    Ensure the code is production-ready, responsive, and follows best practices for accessibility and performance.
    
    For 3D components, use appropriate CSS 3D transforms or consider integrating with Three.js.
    
    Important: Ensure the component fits within a 600x400 pixel container for preview purposes.
    
    Format your response as follows:
    
    HTML:
    ```html
    <!-- Your HTML code here -->
    ```
    
    {css_framework}:
    ```{'css' if css_framework != 'Tailwind CSS' else 'html'}
    /* Your styling code here */
    ```
    
    Ensure the code aligns perfectly with the requested style, functionality, and chosen CSS framework.
    For Tailwind CSS, include classes directly in the HTML.
    Do not include any explanations or comments outside the code blocks.
    """
    
    completion = llm.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model=model_type,
        temperature=0.4
    )
    
    return completion.choices[0].message.content

def extract_code(generated_code, language):
    pattern = rf"```{language}(.*?)```"
    match = re.search(pattern, generated_code, re.DOTALL)
    return match.group(1).strip() if match else ""

def render_component(html_code, css_code, css_framework):
    if css_framework == "Tailwind CSS":
        combined_code = f"""
        <script src="https://cdn.tailwindcss.com"></script>
        <div style="width: 600px; height: 400px; overflow: auto;">
            {html_code}
        </div>
        """
    else:
        combined_code = f"""
        <style>{css_code}</style>
        <div style="width: 600px; height: 400px; overflow: auto;">
            {html_code}
        </div>
        """
    return combined_code

st.set_page_config(layout="wide")

st.title("Advanced Frontend Component Generator")
st.write("By Namah AI")

col1, col2 = st.columns([1, 2])

with col1:
    st.header("Component Settings")
    llm_type = st.selectbox("Select Model", MODELS)
    style = st.selectbox("Select Style", STYLES)
    css_framework = st.selectbox("Select CSS Framework", CSS_FRAMEWORKS)
    
    user_req = st.text_area("Describe your component", height=100, 
                            placeholder="e.g., Create a 3D rotating cube button with hover effects")

    if 'generated_html' not in st.session_state:
        st.session_state.generated_html = ""
    if 'generated_css' not in st.session_state:
        st.session_state.generated_css = ""

    if user_req:
        if st.button("Generate Component"):
            with st.spinner("Generating component..."):
                generated_code = generate_component(user_req, llm_type, style, css_framework)
                st.session_state.generated_html = extract_code(generated_code, "html")
                st.session_state.generated_css = extract_code(generated_code, 'css' if css_framework != 'Tailwind CSS' else 'html')
            st.success("Component generated successfully!")

    if st.button("Regenerate"):
        with st.spinner("Regenerating component..."):
            generated_code = generate_component(user_req, llm_type, style, css_framework)
            st.session_state.generated_html = extract_code(generated_code, "html")
            st.session_state.generated_css = extract_code(generated_code, 'css' if css_framework != 'Tailwind CSS' else 'html')
        st.success("Component regenerated successfully!")
        st.rerun()

with col2:
    if st.session_state.generated_html and st.session_state.generated_css:
        st.header("Generated Component")
        
        tabs = st.tabs(["Preview", "HTML", f"{css_framework}"])
        
        with tabs[0]:
            st.subheader("Live Preview")
            st.components.v1.html(render_component(st.session_state.generated_html, st.session_state.generated_css, css_framework), height=450)
        
        with tabs[1]:
            st.subheader("Edit HTML")
            edited_html = st_ace(value=st.session_state.generated_html, language="html", theme="github", key="html_editor", height=400)
        
        with tabs[2]:
            st.subheader(f"Edit {css_framework}")
            edited_css = st_ace(value=st.session_state.generated_css, 
                                language="css" if css_framework != "Tailwind CSS" else "html", 
                                theme="github", key="css_editor", height=400)
        
        if st.button("Update Live Preview"):
            st.session_state.generated_html = edited_html
            st.session_state.generated_css = edited_css
            st.rerun()

st.sidebar.header("About")
st.sidebar.info("""
This app generates frontend components based on your description. 
You can choose different styles and CSS frameworks, and see a live preview of the result.
You can also edit the generated code and update the preview in real-time.
""")

st.sidebar.header("Tips")
st.sidebar.info("""
1. Be specific in your component description.
2. Try different styles and frameworks for variety.
3. Use the 'Regenerate' button for new ideas.
4. Edit the code to fine-tune the component.
5. The CSS will update automatically when you change the CSS framework.
""")