import os
import streamlit as st
from groq import Groq

# Optional: low-level input libraries
import time
try:
    import keyboard
    import mouse
except:
    keyboard = None
    mouse = None


# --- LLM Macro Generator ---
def generate_macro(task_description, language, groq_client, model_name):
    prompt = f"""
    You are an assistant that generates automation scripts.
    The user will describe a task in natural language.
    Return ONLY the {language} code, nothing else.

    Task: "{task_description}"
    """
    try:
        response = groq_client.chat.completions.create(
            model=model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
            max_tokens=800
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"Error: {str(e)}"


# --- Streamlit App ---
def main():
    # Custom CSS to left-align main content
    st.markdown(
        """
        <style>
            .block-container {
                padding-left: 2rem;
                padding-right: 1rem;
                max-width: 100%;
            }
        </style>
        """,
        unsafe_allow_html=True
    )

    st.title("⚡ AI Macro App")
    st.markdown("""
    Choose between **AI-generated macros** (via natural language)  
    or **manual recording** of keystrokes and clicks.
    """)
    st.sidebar.image("/app/src/2.png")
    st.sidebar.image("/app/src/1.PNG")
    st.info("build by dw")

    # Load Groq API Key
    groq_api_key = "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"
    if not groq_api_key:
        st.error("Missing GROQ_API_KEY in environment or Streamlit secrets.")
        st.stop()

    groq_client = Groq(api_key=groq_api_key)

    # Sidebar model selection
    model_name = st.sidebar.selectbox(
        "Select LLM Model:",
        ["deepseek-r1-distill-llama-70b", "openai/gpt-oss-120b"]
    )

    # Mode selection
    mode = st.sidebar.selectbox("Choose Macro Mode:", ["Generate Macro (AI)", "Record Macro (Manual)"])

    # --- AI Macro Mode ---
    if mode == "Generate Macro (AI)":
        task = st.text_area("📝 Describe your macro task")
        lang = st.sidebar.selectbox("Select output language",
                                    ["Python", "AutoHotkey", "PowerShell", "Macro Recorder (.mrf)"])

        if st.button("Generate Macro"):
            if not task:
                st.error("Please enter a task description.")
                return

            # If Macro Recorder, generate AHK code as base
            actual_lang = "AutoHotkey" if "Macro Recorder" in lang else lang

            with st.spinner(f"Generating script using {model_name}..."):
                script = generate_macro(task, actual_lang, groq_client, model_name)

            # Handle file extensions and wrapping
            if lang == "Macro Recorder (.mrf)":
                script = f"""# Macro Recorder File v5.0
# Task: {task}
# Generated by AI Macro App
BEGIN_MACRO
{script}
END_MACRO
"""
                file_ext = "mrf"
            elif lang == "Python":
                file_ext = "py"
            elif lang == "AutoHotkey":
                file_ext = "ahk"
            elif lang == "PowerShell":
                file_ext = "ps1"
            else:
                file_ext = "txt"

            st.subheader("⚙️ Generated Script")
            st.code(script, language="autohotkey" if file_ext == "ahk" else "python")

            st.download_button("💾 Download Script",
                               script,
                               file_name=f"macro.{file_ext}")

    # --- Record Macro Mode ---
    elif mode == "Record Macro (Manual)":
        st.info("📹 In Record mode, your keystrokes and clicks are logged into a script.")

        if keyboard is None or mouse is None:
            st.error("⚠️ Please install `keyboard` and `mouse` packages to use recording.")
            return

        if "recording" not in st.session_state:
            st.session_state.recording = False
            st.session_state.events = []

        if st.button("Start Recording"):
            st.session_state.recording = True
            st.session_state.events = []
            st.success("Recording started... Perform some actions!")

        if st.button("Stop Recording"):
            st.session_state.recording = False
            st.success("Recording stopped!")

            # Convert to Python (pyautogui-style) script
            lines = ["import pyautogui, time", "time.sleep(2)  # delay before playback"]
            for ev in st.session_state.events:
                if ev["type"] == "key":
                    lines.append(f'pyautogui.typewrite("{ev["value"]}")')
                elif ev["type"] == "click":
                    x, y = ev["value"]
                    lines.append(f"pyautogui.click({x}, {y})")

            script = "\n".join(lines)

            st.subheader("⚙️ Recorded Macro")
            st.code(script, language="python")

            st.download_button("💾 Download Recorded Macro",
                               script,
                               file_name="recorded_macro.py")

        # Simulated event logging
        if st.session_state.recording:
            # Fake log example (in real app, keyboard/mouse hooks would fill this)
            st.session_state.events.append({"type": "click", "value": (200, 300)})
            st.session_state.events.append({"type": "key", "value": "Hello World"})
            st.info("✅ Example events logged (click + keystroke).")


if __name__ == "__main__":
    main()