app.py

import requests
import streamlit as st

# 1. Set page configuration (Must be at the absolute top)
st.set_page_config(layout="wide", page_title="Engineering Document Intelligence System using RAG", page_icon="⚙️")

from config.config import RAW_DIR, OLLAMA_BASE_URL, UPLOADED_DOCS_DIR
from src.vectorstore.embeddings import LocalEmbeddingManager
from src.vectorstore.vector_store import QdrantVectorStoreManager
from src.retrieval.retrieval import HybridRetriever
from src.retrieval.generator import OllamaChatGenerator
from src.retrieval.citation_engine import CitationEngine
from src.ingestion.ingestion_queue import IngestionTaskQueue
from src.monitoring.hallucination_guard import HallucinationComplianceGuard

# Modular UI imports
from src.ui.ui_helpers import safe_markdown, check_ollama_status
from src.ui.animations import get_custom_css
from src.ui.sidebar import render_sidebar
from src.ui.chat_panel import render_chat_panel
from src.ui.pdf_viewer import render_pdf_viewer
from src.ui.engineering_tables import render_table_intelligence
from src.ui.system_monitor import render_system_monitor
from src.ui.ingestion_dashboard import render_ingestion_dashboard

# Inject Custom HUD styling
safe_markdown(get_custom_css())

# --- Resource Caching ---
@st.cache_resource
def get_embedding_manager():
    return LocalEmbeddingManager()

@st.cache_resource
def get_vector_store():
    return QdrantVectorStoreManager(vector_dim=384)

@st.cache_resource
def get_retriever():
    return HybridRetriever(get_vector_store(), get_embedding_manager())

@st.cache_resource
def get_generator():
    return OllamaChatGenerator()

@st.cache_resource
def get_citation_engine():
    return CitationEngine()

@st.cache_resource
def get_ingestion_queue():
    return IngestionTaskQueue(vstore=get_vector_store(), embedding_manager=get_embedding_manager())

@st.cache_resource
def get_hallucination_guard():
    return HallucinationComplianceGuard()

def get_ollama_models():
    """Fetch model registry list from local Ollama host."""
    try:
        res = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=1.5)
        if res.status_code == 200:
            models = [m["name"] for m in res.json().get("models", [])]
            preferred = ["qwen2.5:7b", "llama3.1:latest", "qwen2.5-coder:7b"]
            models_sorted = [m for m in preferred if m in models] + [m for m in models if m not in preferred]
            return models_sorted if models_sorted else preferred
    except Exception:
        pass
    return ["qwen2.5:7b", "llama3.1:latest"]

# --- State Management ---
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "selected_citation" not in st.session_state:
    st.session_state.selected_citation = None
if "last_response_hits" not in st.session_state:
    st.session_state.last_response_hits = []
if "last_query" not in st.session_state:
    st.session_state.last_query = ""
if "last_answer" not in st.session_state:
    st.session_state.last_answer = ""
if "compliance_report" not in st.session_state:
    st.session_state.compliance_report = None
if "retrieval_latency" not in st.session_state:
    st.session_state.retrieval_latency = None
if "llm_latency" not in st.session_state:
    st.session_state.llm_latency = None
if "last_table_result" not in st.session_state:
    st.session_state.last_table_result = None

# --- Main Engine Setup ---
vstore = get_vector_store()
queue_mgr = get_ingestion_queue()
retriever = get_retriever()
generator = get_generator()
guard = get_hallucination_guard()
cit_engine = get_citation_engine()

# --- Automatic Startup Ingestion Migration ---
# Copy any PDFs from data/raw/ to uploaded_documents/ if they are not already in the manifest
import shutil
if RAW_DIR.exists():
    for pdf_path in RAW_DIR.glob("*.pdf"):
        try:
            file_hash = queue_mgr.get_file_hash(str(pdf_path))
            if file_hash not in queue_mgr.manifest:
                dest_file = UPLOADED_DOCS_DIR / pdf_path.name
                if not dest_file.exists():
                    shutil.copy2(str(pdf_path), str(dest_file))
                    print(f"[Migration] Copied '{pdf_path.name}' to uploaded_documents/ for automatic indexing.")
        except Exception as e:
            print(f"[Migration] Error checking/copying '{pdf_path.name}': {e}")

# --- Render Commands Sidebar ---
sidebar_state = render_sidebar(vstore, queue_mgr, get_ollama_models)
selected_model = sidebar_state["selected_model"]
filters = sidebar_state["filters"]

# Sync model states dynamically (model is now passed explicitly to chat functions)
st.session_state.selected_model = selected_model

# --- Header Section ---
st.markdown("""
<div style="margin-bottom:32px;">
    <h1 style='font-family:Orbitron,sans-serif;color:#0EA5E9;font-weight:700;font-size:1.35rem;letter-spacing:-0.2px;margin-bottom:4px;'>🏭 AI-Powered Engineering Document Intelligence System using RAG</h1>
    <p style='color:#64748B;font-size:0.9rem;margin:0;font-family:var(--body-font);line-height:1.4;'>Secure offline technical compliance auditor and specification tracker</p>
</div>
""", unsafe_allow_html=True)

# --- Render Tabbed Workspace ---
tab1, tab2 = st.tabs([
    "💬 Engineering Chat", 
    "📡 System Status"
])

# --- Tab 1: Engineering Chat ---
with tab1:
    render_chat_panel(retriever, generator, guard, filters, selected_model, cit_engine)

# --- Tab 2: System Status & Ingestion Telemetry ---
with tab2:
    render_system_monitor(vstore, queue_mgr, st.session_state.retrieval_latency, st.session_state.llm_latency)
    st.markdown("<hr style='margin:30px 0;border-color:rgba(255,255,255,0.05);'>", unsafe_allow_html=True)
    render_ingestion_dashboard(queue_mgr)