import streamlit as st
import os
from groq import Groq
import concurrent.futures

# --- CONFIGURATION ---
st.set_page_config(page_title="AI Council", page_icon="⚖️", layout="wide")

st.markdown("""
<style>
    .stTextArea textarea { font-size: 16px; }
    .stButton button { height: 50px; font-size: 18px; }
    div[data-testid="stExpander"] details { border: 1px solid #ddd; border-radius: 8px; }
</style>
""", unsafe_allow_html=True)

# --- LOGIN ---
if "authenticated" not in st.session_state:
    st.session_state.authenticated = False

def check_password():
    stored_password = os.environ.get("FRIEND_PASS")
    if not stored_password:
        st.session_state.authenticated = True
        return
    if st.session_state.password_input == stored_password:
        st.session_state.authenticated = True
    else:
        st.error("Wrong password!")

if not st.session_state.authenticated:
    st.title("🔒 Council Access")
    st.text_input("Enter Password", type="password", key="password_input", on_change=check_password)
    st.stop()

# --- CONNECT DIRECTLY TO GROQ ---
# We use the GROQ_API_KEY, not OpenRouter
GROQ_KEY = os.environ.get("GROQ_API_KEY")

if not GROQ_KEY:
    st.error("🚨 Error: Please add GROQ_API_KEY in Settings > Secrets")
    st.stop()

client = Groq(api_key=GROQ_KEY)

# --- THE SURVIVORS (Groq Direct IDs) ---
MEMBERS = [
    # 1. The Smartest (Meta Llama 3.3)
    {"name": "Llama 3.3 70B", "id": "llama-3.3-70b-versatile"},
    
    # 2. The Fast One (Meta Llama 3.1)
    {"name": "Llama 3.1 8B", "id": "llama-3.1-8b-instant"},
]

CHAIRMAN_ID = "llama-3.3-70b-versatile"

def query_groq(model_id, prompt):
    try:
        completion = client.chat.completions.create(
            model=model_id,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=2048,
        )
        return completion.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

# --- INTERFACE ---
st.title("⚖️ The AI Council")
st.markdown("Powered by **Llama 3.3** & **Llama 3.1** (Groq Direct)")

user_query = st.text_area("What is your question?")
start_btn = st.button("Start Debate", type="primary")

if start_btn and user_query:
    
    # STAGE 1
    st.info("🧠 Stage 1: Gathering Opinions...")
    results = {}
    cols = st.columns(len(MEMBERS))
    
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_to_member = {executor.submit(query_groq, m["id"], user_query): m for m in MEMBERS}
        
        for future in concurrent.futures.as_completed(future_to_member):
            member = future_to_member[future]
            ans = future.result()
            results[member["name"]] = ans
            
            idx = MEMBERS.index(member)
            with cols[idx]:
                st.success(member["name"])
                with st.expander("Read Answer", expanded=True):
                    st.write(ans)

    # STAGE 2
    st.warning("⚔️ Stage 2: Peer Review...")
    combined_text = "\n\n".join([f"OPINION FROM {k}:\n{v}" for k, v in results.items()])
    critique_prompt = f"Question: {user_query}\nAnswers:\n{combined_text}\nCritique these answers. Which is best? Be harsh."
    
    critique = query_groq(CHAIRMAN_ID, critique_prompt)
    with st.expander("Read Debate Minutes"):
        st.write(critique)

    # STAGE 3
    st.success("🏆 Stage 3: Final Verdict")
    final_prompt = f"Question: {user_query}\nCritique: {critique}\nSynthesize the perfect answer."
    final_answer = query_groq(CHAIRMAN_ID, final_prompt)
    st.markdown("### Final Answer")
    st.markdown(final_answer)