import streamlit as st
from openai import OpenAI
import os

# --- CONFIGURATION ---
CHOSEN_BREED = "Golden Retrievers"

# Page setup
st.set_page_config(page_title=f"The {CHOSEN_BREED} Expert", page_icon="🐕")
st.title(f"🐕 Only {CHOSEN_BREED}")
st.caption(f"Powered by Perplexity API. Strictly restricted to {CHOSEN_BREED}.")

# --- AUTHENTICATION ---
perplexity_api_key = os.getenv("PERPLEXITY_API_KEY")

if not perplexity_api_key:
    st.error("⚠️ PERPLEXITY_API_KEY not found in secrets.")
    st.stop()

client = OpenAI(
    api_key=perplexity_api_key,
    base_url="https://api.perplexity.ai"
)

# --- CHAT LOGIC ---
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "assistant", "content": f"Hello! I am an expert on {CHOSEN_BREED}. Ask me anything about them!"}
    ]

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

if user_input := st.chat_input():
    
    # 1. Clean up previous failed user messages
    if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
        st.session_state.messages.pop()

    # 2. Display User Message
    st.session_state.messages.append({"role": "user", "content": user_input})
    st.chat_message("user").write(user_input)

    # 3. PREPARE CONTEXT FOR THE AI
    # We grab the last 2 exchanges so the AI knows what "this" or "it" refers to.
    conversation_context = ""
    start_index = max(0, len(st.session_state.messages) - 5) # Get last 5 messages max
    for msg in st.session_state.messages[start_index:]:
        conversation_context += f"{msg['role'].upper()}: {msg['content']}\n"

    # 4. CONSTRUCT THE SMART STRICT PROMPT
    # We give the AI the context so it can resolve pronouns.
    strict_content = f"""
    You are a specialized assistant for {CHOSEN_BREED}.
    
    Here is the recent conversation history for context:
    ---
    {conversation_context}
    ---
    
    Your Task: Respond to the "USER" based *strictly* on the following logic:
    
    1. ANALYZE: Does the USER's latest message relate to {CHOSEN_BREED}? 
       (Check the context! If they say "tell me more" or "how do they eat", assume they mean {CHOSEN_BREED}).
       
    2. IF YES (Valid Topic):
       - Answer the question helpfuly and accurately.
       - DO NOT explain that you are validating the query. Just give the answer.
       
    3. IF NO (Invalid Topic):
       - If the user asks about cats, cars, weather, or other dog breeds:
       - YOU MUST OUTPUT EXACTLY: "REFUSAL_TRIGGERED"
    """

    # 5. Prepare Messages (We strip the history here because we manually injected it into the prompt above)
    api_messages = [
        {"role": "system", "content": "You are a helpful assistant tailored to a specific topic."},
        {"role": "user", "content": strict_content}
    ]

    # 6. Call API
    try:
        response = client.chat.completions.create(
            model="sonar", 
            messages=api_messages,
            temperature=0.1, 
        )
        
        bot_reply = response.choices[0].message.content.strip()

        # 7. POST-PROCESSING GUARDRAIL
        if "REFUSAL_TRIGGERED" in bot_reply:
             final_reply = f"🚫 I cannot answer that. I am strictly programmed to discuss **{CHOSEN_BREED}** only."
        else:
             final_reply = bot_reply

        # 8. Show Result
        st.session_state.messages.append({"role": "assistant", "content": final_reply})
        st.chat_message("assistant").write(final_reply)

    except Exception as e:
        st.error(f"An error occurred: {e}")