import os
import streamlit as st
from groq import Groq
import re

api_key=os.getenv("GROQ_KEY")
client = Groq(api_key=api_key)

# Define available models
MODELS = [
      "llama-3.3-70b-versatile",
    ]

# Configure Streamlit page
st.set_page_config(page_title="Frontend Component Generator", layout="wide")
st.title("Frontend Component Generator")

# Sidebar for settings
with st.sidebar:
    st.header("Settings")
    llm_type = st.selectbox("Select Model", MODELS)
    theme = st.selectbox(
        " Theme",
        [
        "Modern",
        "Vintage",
        "Minimalist",
        "Futuristic",
        "Retro",
        "Glassmorphism",
        "Neumorphism",
        "Material Design",
        "Flat Design",
        "Cyberpunk",
        ],
    )
    st.markdown("---")
    st.markdown("""
    ### How to use:
    1. Enter a description of the component you want.
    2. Click 'Generate Component'.
    3. Preview the component.
    4. Click 'Convert to TSX' to generate a React component.
    5. Download the TSX version.
    6. Click 'Regenerate' for new versions.
    """)

def sanitize_code(code: str) -> str:
    """
    Removes Markdown code fences and any leading/trailing whitespace from the code.
    """
    # Remove opening code fence (``` or ```javascript or ```html)
    code = re.sub(r'^```(?:\w+)?\n?', '', code)
    # Remove closing code fence (```)
    code = re.sub(r'\n?```$', '', code)
    # Strip leading and trailing whitespace
    return code.strip()

def generate_html_js(prompt: str, theme: str):
    """
    Generates HTML, TailwindCSS, and JavaScript code using the GROQ API.
    """
    system_prompt = f"""You are an expert frontend developer.
Create a complete, functional component based on the user's description.

Requirements:
1. Use pure HTML, TailwindCSS, and JavaScript without relying on any frameworks.
2. Utilize the {theme}  theme from TailwindCSS.
3. Ensure the component is responsive.
4. Include form validations if necessary.
5. Add comments to explain the code where appropriate.

**Provide only the raw HTML, TailwindCSS classes, and JavaScript code without any markdown, explanations, or additional text.**"""

    full_prompt = f"Create a component for: {prompt}"

    try:
        completion = client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": full_prompt}
            ],
            model=llm_type,
            temperature=0.4
        )

        content = completion.choices[0].message.content.strip()
        sanitized_content = sanitize_code(content)
        return sanitized_content

    except Exception as e:
        st.error(f"Error generating component: {e}")
        return None

def generate_tsx(html_code: str, js_code: str, theme: str):
    """
    Converts HTML, TailwindCSS, and JavaScript code into a React TSX component using the GROQ API.
    """
    system_prompt = f"""You are an expert React/TypeScript developer.
Convert the following HTML, TailwindCSS, and JavaScript code into a React TypeScript (TSX) component.

Requirements:
1. Use functional components with React Hooks.
2. Incorporate TailwindCSS for styling.
3. Ensure all functionalities from the original JavaScript are preserved.
4. Add appropriate TypeScript types and interfaces.
5. Include necessary imports (e.g., React).
6. Ensure the component is named appropriately and exported as default.

**Provide only the raw TSX component code without any markdown, explanations, or additional text.**"""

    # Combine HTML and JS code into a single prompt
    full_prompt = f"Convert the following code to a React TSX component:\n\nHTML:\n{html_code}\n\nJavaScript:\n{js_code}"

    try:
        completion = client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": full_prompt}
            ],
            model=llm_type,
            temperature=0.4
        )

        tsx_content = completion.choices[0].message.content.strip()
        tsx_content = sanitize_code(tsx_content)
        return tsx_content

    except Exception as e:
        st.error(f"Error converting to TSX: {e}")
        return None

def create_html_preview(html_code: str, js_code: str) -> str:
    """
    Embeds the generated HTML, TailwindCSS, and JavaScript into an HTML template for rendering.
    """
    return f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <script src="https://cdn.tailwindcss.com"></script>
        <style>
            body {{ background-color: #f3f4f6; }}
            #root {{ padding: 20px; }}
        </style>
    </head>
    <body>
        <div id="root">
            {html_code}
        </div>
        <script>
            {js_code}
        </script>
    </body>
    </html>
    """

# Main interface
user_input = st.text_area(
    "Describe the component you want (e.g., 'a login page with email and password fields, remember me checkbox, and a submit button')",
    height=150
)

# Initialize session state
if 'html_code' not in st.session_state:
    st.session_state.html_code = None
if 'js_code' not in st.session_state:
    st.session_state.js_code = None
if 'tsx_code' not in st.session_state:
    st.session_state.tsx_code = None

if user_input:
    generate_button = st.button("Generate Component")

    if generate_button or not st.session_state.html_code:
        with st.spinner("Generating component..."):
            # Generate HTML and JS code
            component_code = generate_html_js(user_input, theme)
            if component_code:
                # Assume that the LLM returns HTML and JS separated by script tags
                # We'll need to parse them
                # For simplicity, let's assume the JS is within <script> tags
                html_part = re.findall(r'([\s\S]*?)<script>', component_code)
                js_part = re.findall(r'<script>([\s\S]*)<\/script>', component_code)

                html_code = html_part[0].strip() if html_part else ""
                js_code = js_part[0].strip() if js_part else ""

                st.session_state.html_code = html_code
                st.session_state.js_code = js_code

                st.success("Component generated successfully!")

    if st.session_state.html_code and st.session_state.js_code:
        # Create tabs for different views
        tab1, tab2 = st.tabs(["Code", "Preview"])

        with tab1:
            st.subheader("Generated HTML")
            st.code(st.session_state.html_code, language="html")

            st.subheader("Generated JavaScript")
            st.code(st.session_state.js_code, language="javascript")

        with tab2:
            st.subheader("Live Preview")
            preview_html = create_html_preview(st.session_state.html_code, st.session_state.js_code)
            st.components.v1.html(preview_html, height=600, scrolling=True)

        # Button to convert to TSX
        convert_button = st.button("Convert to TSX")

        if convert_button and not st.session_state.tsx_code:
            with st.spinner("Converting to TSX..."):
                tsx_code = generate_tsx(st.session_state.html_code, st.session_state.js_code, theme)
                if tsx_code:
                    st.session_state.tsx_code = tsx_code
                    st.success("Conversion to TSX successful!")

        if st.session_state.tsx_code:
            # Create tabs for TSX code and download
            tsx_tab1, tsx_tab2 = st.tabs(["TSX Code", "Download"])

            with tsx_tab1:
                st.subheader("Generated TSX Component")
                st.code(st.session_state.tsx_code, language="typescript")

            with tsx_tab2:
                st.subheader("Download TSX Component")
                st.download_button(
                    label="Download TSX Component",
                    data=st.session_state.tsx_code,
                    file_name="Component.tsx",
                    mime="text/typescript"
                )

        # Regenerate button
        if st.button("Regenerate"):
            # Reset session state
            st.session_state.html_code = None
            st.session_state.js_code = None
            st.session_state.tsx_code = None
            st.rerun()
    else:
        st.warning("Component not parsed, rerun")