import os
import sys
import streamlit as st
from pathlib import Path

# Add project root to path for local execution
sys.path.append(str(Path(__file__).parent.parent))

import psutil
from textwrap import dedent

from src.rag_engine import SatelliteRAG
from src.config import settings

def check_system_health():
    """Check memory availability before processing."""
    process = psutil.Process(os.getpid())
    mem_info = process.memory_info()
    mem_mb = mem_info.rss / 1024 / 1024
    
    # Hugging Face Free Tier ~16GB, Streamlit Cloud ~1GB
    # If using > 80% (just a heuristics check, not hard limit)
    # real check is total memory available
    total_mem = psutil.virtual_memory().total / 1024 / 1024
    avail_mem = psutil.virtual_memory().available / 1024 / 1024
    
    return {
        "used_mb": mem_mb,
        "available_mb": avail_mem,
        "warning": avail_mem < 250 # Warn if < 250MB available
    }

# Setup Page
st.set_page_config(page_title="Space Satellite Assistant", page_icon="🛰️")
st.title("🛰️ Space Satellite Assistant")

# Sidebar
with st.sidebar:
    st.markdown("### Configuration")
    if not settings.GROQ_API_KEY:
        st.error("⚠️ GROQ_API_KEY not found in .env!")
    else:
        st.success("System Online")
    
    st.markdown("---")
    st.markdown("**Engine:** SatelliteRAG v1.0")
    st.markdown("**Knowledge Base:** 637 Satellites")

# Initialize Resources (Cached at startup)
@st.cache_resource
def get_engine():
    try:
        with st.spinner("Initializing RAG Engine (Loading Models & DB)..."):
            return SatelliteRAG()
    except Exception as e:
        # We catch strictly to prevent early app crash
        return f"ERROR: {str(e)}"

# Early Engine Check
engine_status = get_engine()

with st.sidebar:
    st.markdown("---")
    st.markdown("### System Status")
    if isinstance(engine_status, str) and "ERROR" in engine_status:
        st.error(f"🔴 Engine: {engine_status}")
    elif engine_status is None:
        st.warning("🟠 Engine: Initializing...")
    else:
        st.success("🟢 Engine: Ready")

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

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

# Handle Input
if prompt := st.chat_input("Ask about any satellite (e.g., 'What is Gaofen 1?')..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        # Use the already initialized engine
        engine = engine_status

        if isinstance(engine, SatelliteRAG):
            with st.spinner("Analyzing satellite data..."):
                try:
                    # Health Check
                    health = check_system_health()
                    if health["warning"]:
                        st.warning(f"⚠️ Low Memory Warning: Only {health['available_mb']:.0f}MB available. Query might be slow.")
                    
                    # Construct Chat History
                    # We need pairs of (User, AI) from session_state.messages
                    # Excluding the current new prompt which is already appended but not part of 'history' yet for this context
                    chat_history = []
                    msgs = st.session_state.messages[:-1] 
                    for i in range(0, len(msgs) - 1, 2):
                        if msgs[i]["role"] == "user" and msgs[i+1]["role"] == "assistant":
                            chat_history.append((msgs[i]["content"], msgs[i+1]["content"]))
                    
                    response, docs = engine.query(prompt, chat_history=chat_history)

                    st.markdown(response)

                    # Show sources in expander
                    with st.expander(f"View Source Context ({len(docs)} chunks)"):
                        for i, doc in enumerate(docs, 1):
                            st.markdown(f"**Source {i}:** {doc.metadata.get('name')}")
                            st.text(doc.page_content[:300] + "...")

                    st.session_state.messages.append({"role": "assistant", "content": response})
                except Exception as e:
                    st.error("⚠️ An error occurred during analysis.")
                    st.markdown(f"**Error Details:** `{str(e)}`")
                    st.info("💡 **Tip:** If this happens frequently, the system might be out of memory (OOM). Try waiting a few seconds and refreshing.")
        else:
            st.error("⚠️ RAG Engine is not ready. Please check the sidebar for status.")