import streamlit as st
import google.generativeai as genai
import urllib.parse
import re

# --- 1. PAGE CONFIGURATION ---
st.set_page_config(page_title="Nexus Flow Pro ⚡", page_icon="🤖", layout="wide")

# Custom Neon & Cyberpunk CSS
st.markdown("""
    <style>
    .stApp { background-color: #0E1117; color: #FFFFFF; }
    .stChatMessage { border-radius: 15px; border: 1px solid #1E1E1E; margin: 10px 0; }
    /* Thinking Box Style */
    .thinking-box { background-color: #1a1c23; border-left: 4px solid #00ffcc; padding: 10px; margin-bottom: 10px; font-style: italic; color: #00ffcc; }
    </style>
    """, unsafe_allow_html=True)

# --- 2. API KEY SETUP ---
api_key = st.secrets.get("GOOGLE_API_KEY")
if api_key:
    genai.configure(api_key=api_key)
else:
    st.warning("⚠️ Sanjeev, 'Settings > Secrets' mein GOOGLE_API_KEY daalna mat bhulna!")
    st.stop()

# --- 3. PRO SYSTEM INSTRUCTIONS ---
instruction = """
You are Nexus Flow AI Pro, the digital avatar of Sanjeev.
- PERSONALITY: Expert in Video Editing (Punch Edit), Python/C++, and SAT Preparation (Goal 1500+).
- REASONING: Har answer se pehle apni logic <thinking> tags mein likho.
- IMAGES: Agar Sanjeev image mange, toh respond karo: [GENERATE_IMAGE: prompt]
- STYLE: Speak in Hinglish. Be a genius friend.
"""

# --- 4. MODEL & SESSION ---
@st.cache_resource
def load_model():
    return genai.GenerativeModel(model_name="gemini-1.5-flash", system_instruction=instruction)

model = load_model()

if "messages" not in st.session_state:
    st.session_state.messages = []
if "chat_session" not in st.session_state:
    st.session_state.chat_session = model.start_chat(history=[])

# --- 5. UI DISPLAY ---
st.title("Nexus Flow Pro 🤖")
st.caption("Sanjeev's Digital Brain | Powered by Gemini 1.5 Flash")

for m in st.session_state.messages:
    with st.chat_message(m["role"]):
        st.markdown(m["content"])
        if "image" in m:
            st.image(m["image"])

# --- 6. USER INPUT & REASONING LOGIC ---
if prompt := st.chat_input("Kaise help karu Sanjeev?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.status("🔍 Nexus Flow is thinking...", expanded=True) as status:
            try:
                response = st.session_state.chat_session.send_message(prompt)
                full_res = response.text
                img_url = None
                
                # Logic A: Thinking Parser
                if "<thinking>" in full_res:
                    parts = full_res.split("</thinking>")
                    thought = parts[0].replace("<thinking>", "").strip()
                    final_text = parts[1].strip()
                    st.markdown(f'<div class="thinking-box"><b>🧠 Logic Tree:</b><br>{thought}</div>', unsafe_allow_html=True)
                else:
                    final_text = full_res

                # Logic B: Image Parser
                if "[GENERATE_IMAGE:" in final_text:
                    match = re.search(r'\[GENERATE_IMAGE:\s*(.*?)\]', final_text)
                    if match:
                        img_prompt = match.group(1).strip()
                        encoded_prompt = urllib.parse.quote(img_prompt)
                        img_url = f"https://image.pollinations.ai/prompt/{encoded_prompt}?width=1024&height=1024&nologo=true"
                        final_text = f"✅ **Visualizing:** {img_prompt}"

                status.update(label="✅ Analysis Complete!", state="complete", expanded=False)
                
            except Exception as e:
                final