import streamlit as st
from sentence_transformers import SentenceTransformer
import chromadb
import time

# 1. System Configurations
st.set_page_config(page_title="InsightStream AI", page_icon="⚡", layout="wide")

# Custom CSS to force a high-end dark theme aesthetic
st.markdown("""
    <style>
    .main .block-container {padding-top: 2rem;}
    .metric-card {background-color: #1e222b; padding: 15px; border-radius: 8px; border: 1px solid #2d3139;}
    </style>
""", unsafe_allow_html=True)

# 2. Optimized Resource Initializations
@st.cache_resource
def load_analytics_engine():
    return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

@st.cache_resource
def init_vector_vault():
    client = chromadb.Client()
    return client.get_or_create_collection(
        name="stream_intelligence_vault",
        metadata={"hnsw:space": "cosine"}
    )

engine = load_analytics_engine()
vault = init_vector_vault()

# Curated high-value corporate intel data
FEED_PRESETS = [
    "Market Alert: High structural inflation signals aggressive interest rate hikes next quarter.",
    "R&D Update: Quantum entanglement framework successfully processing multi-dimensional data arrays.",
    "Operations Briefing: Deep learning algorithms integrated into production pipelines to automate code audits.",
    "Infrastructure Notice: Storage systems transitioning entirely to vectorized database mathematical modeling.",
    "Wellness Policy: New corporate initiatives mandate daily physical breaks and bio-balanced nutrition tracking.",
    "Logistics Report: Autonomous supply chain mechanisms deployed across regional fulfillment hubs.",
    "Biotech News: Synthetic neural networks showing unprecedented mastery over biochemical sequencing tasks."
]

# 3. Main Dashboard Header Layout
t1, t2 = st.columns([7, 3])
with t1:
    st.title("⚡ INSIGHTSTREAM")
    st.markdown("### *Real-Time Neural Text Processing & Stream Auditing Hub*")
with t2:
    st.write("")
    st.write("")
    # Dynamic KPI metric cards
    st.markdown(
        f'<div class="metric-card">'
        f'📶 <b>SYSTEM STATUS:</b> <span style="color:#00ffcc;">ACTIVE</span><br>'
        f'📦 <b>INDEXED DATA STREAMS:</b> {vault.count()}'
        f'</div>', 
        unsafe_allow_html=True
    )

st.divider()

# 4. Tabbed Workplace Navigation (Eliminating the traditional sidebar)
tab_monitor, tab_ingest = st.tabs(["📊 INTELLIGENCE STREAM MONITOR", "📥 DATA INGESTION ENGINE"])

# --- TAB 1: SEARCH & MONITOR ---
with tab_monitor:
    if vault.count() == 0:
        st.error("🚨 Neural Index Empty: No data feeds detected in the pipeline.")
        if st.button("⚡ Trigger Auto-Seed Pipeline"):
            with st.spinner("Processing framework initialization vectors..."):
                embeddings = engine.encode(FEED_PRESETS).tolist()
                ids = [f"feed_id_{i}" for i in range(len(FEED_PRESETS))]
                vault.add(embeddings=embeddings, documents=FEED_PRESETS, ids=ids)
            st.success("System auto-seeded successfully!")
            st.rerun()
    else:
        # Search controls integrated horizontally into the main screen layout
        ctrl1, ctrl2 = st.columns([3, 1])
        with ctrl1:
            intel_query = st.text_input("🎯 Filter streams via conceptual query:", placeholder="Type a concept, e.g., corporate financial stress or machine learning automation...")
        with ctrl2:
            max_nodes = st.select_slider("Max Nodes to Map:", options=[1, 2, 3, 4, 5], value=3)

        if intel_query:
            with st.spinner("Executing high-dimensional vector search..."):
                t_start = time.time()
                query_vec = engine.encode([intel_query]).tolist()
                payload = vault.query(query_embeddings=query_vec, n_results=max_nodes)
                latency = (time.time() - t_start) * 1000

            st.markdown(f"##### 🛰️ Stream Analysis Complete (`{latency:.2f}ms latency`)")
            
            if payload and payload['documents']:
                feeds = payload['documents'][0]
                metrics = payload['distances'][0]
                node_ids = payload['ids'][0]
                
                # Render clean metric-driven grid modules
                for i in range(len(feeds)):
                    confidence = (1 - metrics[i]) * 100
                    
                    with st.expander(f"🔴 STREAM NODE {node_ids[i]} — Match Confidence: {confidence:.2f}%", expanded=True):
                        st.write(f"**Data Payload:**")
                        st.code(feeds[i], language="markdown")
            else:
                st.warning("No nodes matched the conceptual threshold.")

# --- TAB 2: DATA INGESTION ---
with tab_ingest:
    st.subheader("📥 Raw Text Ingestion Port")
    st.markdown("Manually inject unstructured text blocks into the vector analytics cluster.")
    
    stream_payload = st.text_area("Raw Document Payload:", placeholder="Paste text logs, reports, or data streams here...")
    
    if st.button("🚀 Commit Payload to Cluster"):
        if stream_payload.strip():
            with st.spinner("Vectorizing input stream..."):
                assigned_id = f"feed_user_{vault.count() + 1}"
                payload_vec = engine.encode([stream_payload]).tolist()
                
                vault.add(
                    embeddings=payload_vec,
                    documents=[stream_payload],
                    ids=[assigned_id]
                )
            st.success(f"Success: Nodes updated. Assigned ID: {assigned_id}")
            st.rerun()
        else:
            st.error("Aborted: Ingestion payload cannot be empty.")