import streamlit as st
import openai
import google.generativeai as genai

# Custom CSS for arena styling
st.markdown("""
    <style>
    .fighter-column {
        border: 2px solid #4a4a4a;
        border-radius: 10px;
        padding: 20px;
        margin: 10px;
        background: #1a1a1a;
        box-shadow: 0 4px 8px 0 rgba(0,0,0,0.2);
    }
    .vote-btn {
        width: 100%;
        margin-top: 10px;
    }
    .arena-title {
        text-align: center;
        font-size: 2.5em !important;
        color: #ffd700 !important;
        margin-bottom: 20px !important;
    }
    </style>
""", unsafe_allow_html=True)

def get_openai_response(api_key, model, prompt):
    client = openai.OpenAI(api_key=api_key)
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

def get_gemini_response(api_key, model, prompt):
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel(model)
    response = model.generate_content(prompt)
    return response.text

def get_deepseek_response(api_key, model, prompt):
    client = openai.OpenAI(
        api_key=api_key,
        base_url="https://api.deepseek.com/v1",
    )
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# Initialize session state for voting
if 'votes' not in st.session_state:
    st.session_state.votes = {'OpenAI': 0, 'Gemini': 0, 'DeepSeek': 0}

# Sidebar configuration
with st.sidebar:
    st.header("⚙️ Battle Configuration")
    st.subheader("API Keys")
    openai_api_key = st.text_input("OpenAI Key", type="password")
    gemini_api_key = st.text_input("Gemini Key", type="password")
    deepseek_api_key = st.text_input("DeepSeek Key", type="password")
    
    st.subheader("Model Selection")
    openai_model = st.text_input("OpenAI Model", value="o1-mini")
    gemini_model = st.text_input("Gemini Model", value="gemini-2.0-flash-exp")
    deepseek_model = st.text_input("DeepSeek Model", value="deepseek-reasoner")

# Main Arena
st.markdown('<p class="arena-title">🤖 LLM BATTLE ARENA 🥊</p>', unsafe_allow_html=True)
prompt = st.text_area("Enter your battle prompt:", height=150, help="The prompt that will determine the AI champion!")

if st.button("🚀 Start Battle!"):
    if not prompt.strip():
        st.error("Please enter a battle prompt!")
    else:
        responses = {}
        
        # Get all responses
        with st.spinner("⚔️ Models are battling it out..."):
            col1, col2, col3 = st.columns(3)
            
            # OpenAI Fighter
            with col1:
                with st.container():
                    st.markdown('<div class="fighter-column">', unsafe_allow_html=True)
                    st.markdown("### 🥊 OpenAI\n`" + openai_model + "`")
                    try:
                        if openai_api_key:
                            response = get_openai_response(openai_api_key, openai_model, prompt)
                            responses["OpenAI"] = response
                            st.markdown("---")
                            st.write(response)
                        else:
                            st.error("Missing API Key!")
                    except Exception as e:
                        st.error(f"Battle error: {str(e)}")
                    st.markdown('</div>', unsafe_allow_html=True)

            # Gemini Fighter
            with col2:
                with st.container():
                    st.markdown('<div class="fighter-column">', unsafe_allow_html=True)
                    st.markdown("### 🥊 Gemini\n`" + gemini_model + "`")
                    try:
                        if gemini_api_key:
                            response = get_gemini_response(gemini_api_key, gemini_model, prompt)
                            responses["Gemini"] = response
                            st.markdown("---")
                            st.write(response)
                        else:
                            st.error("Missing API Key!")
                    except Exception as e:
                        st.error(f"Battle error: {str(e)}")
                    st.markdown('</div>', unsafe_allow_html=True)

            # DeepSeek Fighter
            with col3:
                with st.container():
                    st.markdown('<div class="fighter-column">', unsafe_allow_html=True)
                    st.markdown("### 🥊 DeepSeek\n`" + deepseek_model + "`")
                    try:
                        if deepseek_api_key:
                            response = get_deepseek_response(deepseek_api_key, deepseek_model, prompt)
                            responses["DeepSeek"] = response
                            st.markdown("---")
                            st.write(response)
                        else:
                            st.error("Missing API Key!")
                    except Exception as e:
                        st.error(f"Battle error: {str(e)}")
                    st.markdown('</div>', unsafe_allow_html=True)

        # Voting Section
        st.markdown("---")
        st.subheader("🏆 Judge the Responses!")
        cols = st.columns(3)
        with cols[0]:
            if st.button("Vote OpenAI 👑", key="vote_openai"):
                st.session_state.votes["OpenAI"] += 1
        with cols[1]:
            if st.button("Vote Gemini 👑", key="vote_gemini"):
                st.session_state.votes["Gemini"] += 1
        with cols[2]:
            if st.button("Vote DeepSeek 👑", key="vote_deepseek"):
                st.session_state.votes["DeepSeek"] += 1

        # Leaderboard
        st.markdown("---")
        st.subheader("📊 Battle Leaderboard")
        sorted_votes = sorted(st.session_state.votes.items(), key=lambda x: x[1], reverse=True)
        
        for i, (model, votes) in enumerate(sorted_votes):
            trophy = "🥇" if i == 0 else "🥈" if i == 1 else "🥉"
            st.markdown(f"{trophy} **{model}**: {votes} votes")
            st.progress(votes / (sum(st.session_state.votes.values()) or 1))

# How to Run reminder in sidebar
with st.sidebar:
    st.markdown("---")
    st.markdown("**How to start the battle:**")
    st.markdown("1. Enter API keys 🔑\n2. Set models 🧠\n3. Write prompt 📝\n4. Click battle button! ⚔️")