import streamlit as st
import openai

# Set Streamlit page configuration
st.set_page_config(page_title="AI Prompt Enhancer", layout="centered")

# OpenAI API Key (Store securely in Streamlit secrets)
OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]

# Initialize OpenAI Client
client = openai.OpenAI(api_key=OPENAI_API_KEY)

# 📌 Dictionary of Prompting Techniques
prompt_techniques = {
    "🔗 Chain of Thought (CoT)": "Breaks down reasoning into step-by-step logic.",
    "📚 Few-Shot Learning": "Adds relevant examples for better AI understanding.",
    "🧩 ReAct (Reasoning + Acting)": "Improves AI decision-making with iterative steps.",
    "🌲 Tree of Thoughts (ToT)": "Generates multiple possible reasoning paths.",
    "💡 Self-Consistency": "Reframes the prompt to generate diverse responses.",
    "📑 HyDE (Hypothetical Doc Embeddings)": "Structures prompts for better retrieval-based outputs.",
    "🔄 Least-to-Most Prompting": "Breaks complex tasks into smaller, manageable steps.",
    "📊 Graph Prompting": "Optimizes structured data queries and relationships."
}

# 🎯 Title
st.title("✨ AI Prompt Enhancer")

# 💡 Prompt Input Box
user_prompt = st.text_area("💡 Prompt Idea Box", placeholder="Enter your prompt here...")

# 📌 Dropdown for Selecting a Prompting Technique
selected_technique = st.selectbox(
    "🛠️ Choose a Prompting Technique",
    options=list(prompt_techniques.keys()),
    format_func=lambda x: f"{x} - {prompt_techniques[x]}"  # Display name + description
)

# 🚀 Enhance Prompt Button
if st.button("Enhance Prompt"):
    if user_prompt.strip() == "":
        st.warning("⚠️ Please enter a prompt before enhancing.")
    else:
        try:
            # Ensure max_tokens is an integer
            # max_tokens = int(min(len(user_prompt) * 1.5, 500))  # Convert to integer
            max_tokens = None

            # Generate enhanced prompt using OpenAI API
            completion = client.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {
                        "role": "system", 
                        "content": f"""
                        Your task is to rewrite and improve the given prompt using the '{selected_technique}' technique.
                        **DO NOT ANSWER** the prompt. 
                        **DO NOT GENERATE ADDITIONAL INFORMATION.** 
                        **ONLY return an improved version of the prompt, making it clearer, more specific, and optimized for AI response generation.**
                        """
                    },
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,  # Low randomness
                max_tokens=max_tokens  # Ensure integer value
            )
            # Ensure OpenAI returns only plain text
            enhanced_prompt = completion.choices[0].message.content.strip()

            # Remove markdown-style bold formatting (**text**)
            import re
            enhanced_prompt = re.sub(r"\*\*(.*?)\*\*", r"\1", enhanced_prompt)

            # 📌 Display Enhanced Prompt in a Larger Text Box
            st.subheader("🔹 Enhanced Prompt:")
            st.text_area(" ", value=enhanced_prompt, height=450 )  # Bigger text box

        except Exception as e:
            st.error(f"❌ Error: {str(e)}")