import streamlit as st
import os
from dotenv import load_dotenv

# --- KAEL'S SUBSYSTEMS ---
from tig_engine import IntelliMod
from intellimod_bridge import IntelliModBridge
from librarian import Librarian

load_dotenv()

# --- PAGE CONFIG ---
st.set_page_config(
    page_title="Kael | Mission Control",
    page_icon="🧬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- CSS FOR POLISH ---
st.markdown("""
<style>
    .stChatInput {position: fixed; bottom: 0; padding-bottom: 1rem; z-index: 100;}
    .block-container {padding-bottom: 5rem;}
</style>
""", unsafe_allow_html=True)

# --- INITIALIZE SYSTEMS ---
@st.cache_resource
def load_brain():
    base_path = "/workspaces/collaborator_agent/memory"
    
    # 1. The Tools
    tig = IntelliMod()            # The Router Engine
    bridge = IntelliModBridge()   # The Registry Reader
    lib = Librarian(base_path)    # The Long-Term Memory
    
    # 2. Load Identity
    profile_path = os.path.join(base_path, "profile_core.md")
    if os.path.exists(profile_path):
        with open(profile_path, "r") as f: identity = f.read()
    else:
        identity = "You are Kael, a collaborative AI agent working with Jaccob."
        
    return tig, bridge, lib, identity

tig, bridge, librarian, core_identity = load_brain()

# --- SIDEBAR: MISSION CONTROL ---
with st.sidebar:
    st.title("🎛️ Control Deck")
    
    # 1. OPERATION MODE
    st.subheader("🎯 Operational Mode")
    mode = st.radio(
        "Select Protocol:",
        ["General Chat", "IntelliMod OS", "Deep Research", "Coding / Dev", "Brainstorming"],
        index=0,
        help="IntelliMod OS activates the Prompt Compiler logic."
    )

    st.divider()

    # 2. ENGINE SELECTOR
    st.subheader("⚙️ Engine Selector")
    engine_choice = st.selectbox(
        "Active Model:",
        ["Auto-Pilot (TIG Router)", "claude-sonnet-4-5-20250929", "gpt-5.1-chat-latest", "gemini-3-pro-preview", "gemini-2.5-flash"],
        index=0
    )
    force_model = None if "Auto" in engine_choice else engine_choice
    
    st.divider()
    
    if st.button("🌙 Save & Sleep"):
        st.success("Memory Synced to Drive.")

# --- CHAT LOGIC ---
st.header(f"Kael Online")
st.caption(f"Connected to: **Jaccob** | Protocol: **{mode}** | Engine: **{engine_choice}**")

if "messages" not in st.session_state:
    st.session_state.messages = []

# Display History
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Handle Input
if prompt := st.chat_input("Direct Kael..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        
        with st.spinner("Processing..."):
            # A. RETRIEVAL
            retrieved_knowledge = librarian.query(prompt, n_results=2)
            
            # B. TIG / INTELLIMOD LOGIC
            intent = tig.detect_intent(prompt)
            
            # Show Visual Card if in IntelliMod Mode
            if mode == "IntelliMod OS" and intent != "chat":
                card_data = bridge.get_tig_recommendation(intent)
                if card_data:
                    with st.status(f"⚡ IntelliMod System: {intent.upper()}", expanded=True):
                        st.write(f"**Selected Card:** `{card_data['card_name']}`")
                        st.write(f"**Reason:** {card_data['category']} allows for optimized {intent}.")
            
            # C. SYSTEM PROMPT ASSEMBLY (The "Prompt Constructor" Logic)
            mode_instructions = "SYSTEM: ACT AS A HELPFUL ASSISTANT." # Default
            
            if mode == "IntelliMod OS":
                mode_instructions = """
                SYSTEM GOAL: YOU ARE THE 'MPI RUNTIME COMPILER'.
                1. ANALYZE the user's request.
                2. SELECT relevant System Cards (SC_) and V-Cards (VC_) from your memory/files.
                3. COMPILE a structured, optimized prompt artifact.
                4. DO NOT just answer the question. OUTPUT the prompt design itself.
                """
            elif mode == "Coding / Dev":
                mode_instructions = "SYSTEM: ACT AS SENIOR SOFTWARE ENGINEER. PRIORITIZE CLEAN CODE."
            elif mode == "Deep Research":
                mode_instructions = "SYSTEM: ACT AS RESEARCH ANALYST. CITE SOURCES."
            elif mode == "Brainstorming":
                mode_instructions = "SYSTEM: ACT AS CREATIVE PARTNER. OFFER DIVERGENT IDEAS."

            # D. FINAL PROMPT CREATION (The Fix: Defined HERE, always)
            full_system_prompt = f"""
            SYSTEM IDENTITY:
            {core_identity}
            
            CURRENT USER: Jaccob
            CURRENT MODE: {mode}
            
            INSTRUCTIONS:
            {mode_instructions}
            
            RELEVANT KNOWLEDGE (From Library):
            {retrieved_knowledge}
            
            USER PROMPT:
            {prompt}
            """
            
            # E. EXECUTION
            response = tig.run_tig_pipeline(full_system_prompt, force_model=force_model)
            
            message_placeholder.markdown(response)

    st.session_state.messages.append({"role": "assistant", "content": response})