# --- 1. SQLITE HACK (Must be top) ---
__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

import streamlit as st
import os
import tempfile
import time
import logging
import warnings
import json
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_cohere import CohereEmbeddings
from agent import get_agent_app
# --- 2. CONFIG & SETUP ---
st.set_page_config(
    page_title="AutoDev Agent", 
    page_icon="🤖", 
    layout="wide",
    initial_sidebar_state="expanded"
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("AutoDev_UI")
warnings.filterwarnings("ignore")


# --- 3. CUSTOM CSS (Professional Dark/Glass Look) ---
st.markdown("""
<style>
    /* Main Background adjustments if needed */
    .stApp {
        background-color: #0e1117;
    }
    
    /* Sidebar Styling */
    [data-testid="stSidebar"] {
        background-color: #161b22;
        border-right: 1px solid #30363d;
    }
    
    /* Chat Message Styling */
    .stChatMessage {
        background-color: #21262d;
        border: 1px solid #30363d;
        border-radius: 10px;
        padding: 15px;
    }
    
    /* User Message distinct color */
    div[data-testid="stChatMessage"]:nth-child(odd) {
        background-color: #1c2128;
        border-left: 4px solid #00D084; /* Green accent */
    }
    
    /* Bot Message distinct color */
    div[data-testid="stChatMessage"]:nth-child(even) {
        background-color: #161b22;
        border-left: 4px solid #2F80ED; /* Blue accent */
    }

    /* Expander Styling for Logs */
    .streamlit-expanderHeader {
        font-family: 'Monospace';
        font-size: 0.9em;
        color: #8b949e;
    }
</style>
""", unsafe_allow_html=True)

# --- SESSION STATE ---
if "messages" not in st.session_state:
    st.session_state.messages = []
if "workflow" not in st.session_state:
    st.session_state.workflow = None

# --- 4. HELPER FUNCTIONS ---

def ingest_pdf(uploaded_file):
    """Ingests PDF into ChromaDB"""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(uploaded_file.getvalue())
            tmp_path = tmp.name
        
        loader = PyPDFLoader(tmp_path)
        docs = loader.load()
        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = splitter.split_documents(docs)
        
        api_key = os.environ.get("COHERE_API_KEY")
        if not api_key: return False, "COHERE_API_KEY missing in .env"
        
        embeddings = CohereEmbeddings(model="embed-english-v3.0", cohere_api_key=api_key)
        # Persistent path fix
        db = Chroma(persist_directory="./chroma_db_groq_v1", embedding_function=embeddings, collection_name="rag_groq_final")
        db.add_documents(splits)
        
        return True, len(splits)
    except Exception as e:
        return False, str(e)

def format_agent_output(content):
    """Cleans up the raw content for display"""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        # Extract content from list of messages
        return "\n".join([msg.content for msg in content if hasattr(msg, 'content')])
    return str(content)

# --- 5. SIDEBAR ---
with st.sidebar:
    st.title("🤖 AutoDev Agent")
    st.caption("v2.0 | Enterprise Edition")
    st.markdown("---")
    
    # PDF Upload
    st.subheader("📂 Knowledge Base")
    uploaded_file = st.file_uploader("Upload CV/Docs", type=["pdf"])
    
    if uploaded_file:
        if st.button("📥 Ingest Document", use_container_width=True):
            with st.spinner("Processing Knowledge Graph..."):
                success, msg = ingest_pdf(uploaded_file)
                if success: 
                    st.success(f"Indexed {msg} chunks successfully!")
                    time.sleep(1)
                    st.rerun()
                else: 
                    st.error(f"Failed: {msg}")
    
    st.markdown("---")
    
    # Status Indicator
    if st.session_state.workflow:
        st.success("🟢 System Online")
    else:
        st.warning("🔴 System Offline")
        
    if st.button("🧹 Clear Chat History", use_container_width=True):
        st.session_state.messages = []
        st.rerun()

# --- 6. MAIN APP LOGIC ---

st.subheader("💬 Career & Tech Assistant")

# Initialize Agent lazily
if not st.session_state.workflow:
    with st.spinner("🚀 Booting up Multi-Agent System..."):
        try:
            st.session_state.workflow = get_agent_app()
            st.rerun()
        except Exception as e:
            st.error(f"Failed to initialize Agent: {e}")
            st.stop()

# Display Chat History
for msg in st.session_state.messages:
    avatar = "🧑‍💻" if msg["role"] == "user" else "🤖"
    with st.chat_message(msg["role"], avatar=avatar):
        st.markdown(msg["content"])

# --- 7. INPUT HANDLING ---
if prompt := st.chat_input("Ask about skills, projects, or upload a doc..."):
    
    # 1. User Message Display
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar="🧑‍💻"):
        st.markdown(prompt)

    # 2. Agent Processing
    with st.chat_message("assistant", avatar="🤖"):
        # Placeholder for final answer
        response_placeholder = st.empty()
        
        # Status container for "Thinking"
        with st.status("🧠 Orchestrating Agents...", expanded=True) as status:
            final_response = "I couldn't generate a response."
            
            try:
                inputs = {"messages": [{"role": "user", "content": prompt}]}
                
                # Streaming Loop
                for event in st.session_state.workflow.stream(inputs):
                    for agent_name, agent_data in event.items():
                        
                        # Extract the actual message content cleanly
                        raw_messages = agent_data.get("messages", [])
                        if raw_messages:
                            last_msg = raw_messages[-1]
                            content = last_msg.content
                            
                            # Determine Icon based on Agent
                            icon = "⚙️"
                            if "graph" in agent_name: icon = "🕸️"
                            elif "vector" in agent_name: icon = "📄"
                            elif "web" in agent_name: icon = "🌐"
                            elif "supervisor" in agent_name: icon = "👔"

                            # 1. Show high-level step
                            st.write(f"**{icon} {agent_name.replace('_', ' ').title()}**")
                            
                            # 2. Show details in dropdown (No ugly JSON)
                            with st.expander(f"View {agent_name} Output", expanded=False):
                                st.markdown(f"```text\n{content}\n```")
                            
                            # Update final response tracking
                            final_response = content

                status.update(label="✅ Task Completed", state="complete", expanded=False)
            
            except Exception as e:
                status.update(label="❌ Error Occurred", state="error")
                st.error(f"Runtime Error: {e}")
        
        # 3. Final Answer Display (Clean)
        response_placeholder.markdown(final_response)
        st.session_state.messages.append({"role": "assistant", "content": final_response})