src/app.py

import streamlit as st
from engine import AdvancedPromptOptimizer
from llm_optimizer import optimize_with_llm, PERSONAS, get_accurate_token_count, optimize_with_agent
from dotenv import load_dotenv
import os

load_dotenv()

cost_model = {
    "GPT-4": (0.01, 0.03),
    "GPT-5": (0.012, 0.04),  # Premium pricing for latest model
    "Claude Opus": (0.015, 0.075),
    "Claude Sonnet": (0.003, 0.015),
    "LLaMA 2": (0.012, 0.04),
    "Custom": (None, None),
}


def format_cost(tokens, cost_per_k):
    return f"${tokens * cost_per_k / 1000:.4f}"

def run_chat_interface():
    """Chat interface for agent-based interactions"""
    st.subheader("💬 Chat with AI Agent")
    
    # Initialize chat history
    if "chat_messages" not in st.session_state:
        st.session_state.chat_messages = []
    
    # Display chat messages
    for message in st.session_state.chat_messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Chat input
    if prompt := st.chat_input("Ask the agent about prompt optimization..."):
        # Add user message to chat history
        st.session_state.chat_messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Generate assistant response
        with st.chat_message("assistant"):
            response = "🤖 Advanced agent functionality coming soon! This will include:\n\n" \
                      "• Memory of conversation context\n" \
                      "• Interactive clarification questions\n" \
                      "• Multi-turn optimization refinement\n" \
                      "• Personalized optimization strategies\n\n" \
                      f"For now, you asked: '{prompt}'"
            st.markdown(response)
        
        # Add assistant response to chat history
        st.session_state.chat_messages.append({"role": "assistant", "content": response})

def main():
    st.set_page_config(layout="wide", page_title="Prompt Optimizer")
    st.title("👋✨ Welcome from the PromptCraft Team ✨💡")
    st.title("🚀 AI Prompt Optimizer")

    col1, col2 = st.columns([0.6, 0.4])  # 60/40 split for better space utilization

    with col1:
        st.subheader("LLM Cost")
        model = st.selectbox("Select LLM Model", list(cost_model.keys()))

        if model == "Custom":
            input_cost = st.number_input("Input Cost ($/1K tokens)", 0.01, 1.0, 0.03)
            output_cost = st.number_input("Output Cost ($/1K tokens)", 0.01, 1.0, 0.06)
        else:
            input_cost, output_cost = cost_model[model]

        st.subheader("Optimization Model")
        
        # Create columns for the optimizer section
        opt_col1, opt_col2 = st.columns([1, 1])
        
        with opt_col1:
            optimizer_model = st.selectbox(
                "Choose Optimizer", 
                [
                    "spaCy + Lemminflect",
                    "GPT-5 (Simple LLM Optimization)", 
                    "GPT-5 Search (Agent-Based with Search)", 
                    "Agent (Chat Interface)"
                ]
            )

        persona = "Default"
        api_key_input = ""
        tavily_api_key_input = ""
        aggressiveness = 1.0  # Initialize default value
        
        # Handle different optimizer configurations
        if optimizer_model in ["GPT-5 (Simple LLM Optimization)", "GPT-5 Search (Agent-Based with Search)"]:
            with opt_col2:
                persona = st.selectbox("Choose Persona", list(PERSONAS.keys()))
            
            # API Keys in the same row
            api_col1, api_col2 = st.columns([1, 1])
            with api_col1:
                api_key_input = st.text_input("AIMLAPI API Key (optional)", type="password", help="If you don't provide a key, the one in your .env file will be used.")
            
            if optimizer_model == "GPT-5 Search (Agent-Based with Search)":
                with api_col2:
                    tavily_api_key_input = st.text_input("Tavily API Key (optional)", type="password", help="If you don't provide a key, the one in your .env file will be used.")
        
        elif optimizer_model == "spaCy + Lemminflect":
            with opt_col2:
                aggressiveness = st.slider(
                    "Optimization Level",
                    0.0,
                    1.0,
                    0.7,
                    help="Higher = more aggressive shortening",
                )
        
        elif optimizer_model == "Agent (Chat Interface)":
            with opt_col2:
                persona = st.selectbox("Choose Persona", list(PERSONAS.keys()))
            
            # API Keys for agent
            api_col1, api_col2 = st.columns([1, 1])
            with api_col1:
                api_key_input = st.text_input("AIMLAPI API Key (optional)", type="password", help="Required for agent functionality")
            with api_col2:
                tavily_api_key_input = st.text_input("Tavily API Key (optional)", type="password", help="Optional for search capabilities")

        # Show different interfaces based on optimizer choice
        if optimizer_model == "Agent (Chat Interface)":
            # Chat interface takes over the main area
            st.markdown("---")
            run_chat_interface()
            return  # Exit early for chat interface
        else:
            # Traditional prompt optimization interface
            prompt = st.text_area(
                "Original Prompt", height=200, placeholder="Paste your AI prompt here..."
            )

    # Optimization button and results (only for non-chat interfaces)
    if st.button("Optimize", type="primary"):
        if optimizer_model == "spaCy + Lemminflect":
            optimizer = AdvancedPromptOptimizer()
            optimized, orig_toks, new_toks = optimizer.optimize(prompt, aggressiveness)
            
        elif optimizer_model == "GPT-5 (Simple LLM Optimization)":
            api_key = api_key_input if api_key_input else os.getenv("AIMLAPI_API_KEY")
            if not api_key or api_key == "<YOUR_API_KEY>":
                st.error("Please set your AIMLAPI_API_KEY in the .env file or enter it above.")
                return
            
            optimized = optimize_with_llm(prompt, api_key, persona)
            
            orig_toks = get_accurate_token_count(prompt, "gpt-5")
            new_toks = get_accurate_token_count(optimized, "gpt-5")
            
        elif optimizer_model == "GPT-5 Search (Agent-Based with Search)":
            api_key = api_key_input if api_key_input else os.getenv("AIMLAPI_API_KEY")
            tavily_api_key = tavily_api_key_input if tavily_api_key_input else os.getenv("TAVILY_API_KEY")
            if not api_key or api_key == "<YOUR_API_KEY>":
                st.error("Please set your AIMLAPI_API_KEY in the .env file or enter it above.")
                return
            
            optimized = optimize_with_agent(prompt, api_key, persona, tavily_api_key=tavily_api_key)
            
            orig_toks = get_accurate_token_count(prompt, "gpt-5")
            new_toks = get_accurate_token_count(optimized, "gpt-5")

        if orig_toks == 0:
            st.warning("Please enter a valid prompt.")
            return

        # Calculate savings
        token_savings = orig_toks - new_toks
        percent_savings = (token_savings / orig_toks) * 100 if orig_toks > 0 else 0
        input_cost_savings = token_savings * input_cost / 1000
        output_cost_savings = token_savings * output_cost / 1000
        total_cost_savings = input_cost_savings + output_cost_savings

        with col1:
            st.subheader("Optimized Prompt")
            st.code(optimized, language="text")

            # Add download button
            st.download_button(
                "📥 Download Optimized Prompt",
                optimized,
                file_name="optimized_prompt.txt",
            )

        with col2:
            st.subheader("💰 Optimization Results")
            
            # Show method-specific info
            if optimizer_model == "GPT-5 Search (Agent-Based with Search)":
                st.info("🔍 Research-Enhanced Optimization with intelligent search integration.")
            elif optimizer_model == "GPT-5 (Simple LLM Optimization)":
                st.info("⚡ Fast LLM-based optimization without search.")
            elif optimizer_model == "spaCy + Lemminflect":
                st.info("🔧 Rule-based linguistic optimization.")

            # Token Savings - Percentage First
            st.markdown(
                f"""
                <div style=\"background-color:#f0f2f6;padding:15px;border-radius:10px;margin-bottom:15px;">
                    <h3 style=\"color:#2e86c1;margin-top:0;\">Token Reduction</h3>
                    <div style=\"font-size:28px;font-weight:bold;color:#27ae60;text-align:center;\">
                        {percent_savings:.1f}%
                    </div>
                    <div style=\"text-align:center;color:#7f8c8d;font-size:14px;\">
                        {token_savings} tokens saved
                    </div>
                </div>
                """,
                unsafe_allow_html=True,
            )

            # Cost Savings - Percentage First
            if orig_toks > 0 and (input_cost + output_cost) > 0:
                cost_percent_savings = (
                    total_cost_savings
                    / (orig_toks * (input_cost + output_cost) / 1000)
                    * 100
                )
            else:
                cost_percent_savings = 0
            st.markdown(
                f"""
                <div style=\"background-color:#f0f2f6;padding:15px;border-radius:10px;margin-bottom:15px;">
                    <h3 style=\"color:#2e86c1;margin-top:0;\">Cost Reduction</h3>
                    <div style=\"font-size:28px;font-weight:bold;color:#27ae60;text-align:center;\">
                        {cost_percent_savings:.1f}%
                    </div>
                    <div style=\"text-align:center;color:#7f8c8d;font-size:14px;\">
                        ${total_cost_savings:.4f} saved per call
                    </div>
                </div>
                """,
                unsafe_allow_html=True,
            )

            # Visual indicator with percentage
            st.progress(min(1.0, max(0.0, percent_savings / 100)))
            st.caption(f"Prompt reduced to {100-percent_savings:.1f}% of original size")
            
            # Detailed Breakdown
            with st.expander("📊 Cost Analysis"):
                col_a, col_b = st.columns(2)
                with col_a:
                    st.markdown(
                        f"**Input Cost**\n\n"
                        f"Original: {format_cost(orig_toks, input_cost)}\n\n"
                        f"Optimized: {format_cost(new_toks, input_cost)}\n\n"
                        f"Saved: {format_cost(token_savings, input_cost)}"
                    )
                with col_b:
                    st.markdown(
                        f"**Output Cost**\n\n"
                        f"Original: {format_cost(orig_toks, output_cost)}\n\n"
                        f"Optimized: {format_cost(new_toks, output_cost)}\n\n"
                        f"Saved: {format_cost(token_savings, output_cost)}"
                    )

            # Optimization report
            with st.expander("🔍 Applied Optimizations"):
                if optimizer_model == "GPT-5 Search (Agent-Based with Search)":
                    st.markdown("🤖 **Agent-Based Optimization**")
                    st.markdown("• Intelligent search query generation")
                    st.markdown("• Context-aware prompt enhancement")
                    st.markdown("• Persona-specific optimization")
                    st.markdown("• Autonomous decision-making process")
                elif optimizer_model == "GPT-5 (Simple LLM Optimization)":
                    st.markdown("⚡ **Simple LLM Optimization**")
                    st.markdown("• Direct prompt optimization")
                    st.markdown("• Persona-specific guidelines")
                    st.markdown("• Fast processing")
                    st.markdown("• No search enhancement")
                else:
                    st.markdown("🔧 **Rule-Based Optimization**")
                    st.markdown(f"• Optimization aggressiveness: {aggressiveness*100:.0f}%")

                st.markdown("### Share Your Savings")
                optimization_method = {
                    "spaCy + Lemminflect": f"Rule-based (Level: {aggressiveness*100:.0f}%)",
                    "GPT-5 (Simple LLM Optimization)": f"Simple LLM ({persona})",
                    "GPT-5 Search (Agent-Based with Search)": f"Agent+Search ({persona})"
                }
                st.code(
                    f"Saved {token_savings} tokens (${total_cost_savings:.4f}) with #PromptOptimizer\n"
                    f"Method: {optimization_method.get(optimizer_model, 'Unknown')}"
                )

if __name__ == "__main__":
    main()