import streamlit as st
import os
import json
import time
from datetime import datetime
from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_classic.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate

# ─── Page Config ────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Multi-Model Research Assistant",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ─── Custom CSS ─────────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600&family=Syne:wght@400;700;800&display=swap');

/* Base */
html, body, [class*="css"] {
    font-family: 'Syne', sans-serif;
}

/* Hide default streamlit branding */
#MainMenu, footer, header {visibility: hidden;}

/* Main background */
.stApp {
    background: #0d0f14;
    color: #e2e8f0;
}

/* Sidebar */
[data-testid="stSidebar"] {
    background: #111318 !important;
    border-right: 1px solid #1e2130;
}
[data-testid="stSidebar"] * {
    color: #c4cad6 !important;
}

/* Top title bar */
.top-bar {
    background: linear-gradient(90deg, #0f1117 0%, #141822 100%);
    border-bottom: 1px solid #1e2535;
    padding: 1rem 2rem;
    margin: -1rem -1rem 2rem -1rem;
    display: flex;
    align-items: center;
    gap: 1rem;
}
.top-bar h1 {
    font-family: 'Syne', sans-serif;
    font-weight: 800;
    font-size: 1.6rem;
    background: linear-gradient(135deg, #60a5fa, #a78bfa, #34d399);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0;
}

/* Model cards */
.model-card {
    background: #13161f;
    border: 1px solid #1e2535;
    border-radius: 12px;
    padding: 1.2rem;
    height: 100%;
    transition: border-color 0.2s;
}
.model-card:hover {
    border-color: #334155;
}
.model-badge {
    display: inline-block;
    padding: 3px 10px;
    border-radius: 20px;
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.72rem;
    font-weight: 600;
    margin-bottom: 0.8rem;
    letter-spacing: 0.03em;
}
.badge-Meta { background: #1a2a3a; color: #60a5fa; border: 1px solid #2563eb44; }
.badge-qwen     { background: #1f2a1a; color: #4ade80; border: 1px solid #16a34a44; }
.badge-gemma    { background: #2a1a2a; color: #c084fc; border: 1px solid #9333ea44; }

.model-answer {
    font-size: 0.9rem;
    line-height: 1.7;
    color: #cbd5e1;
    min-height: 120px;
}

/* Chat history item */
.history-item {
    background: #13161f;
    border: 1px solid #1e2535;
    border-radius: 8px;
    padding: 0.7rem 1rem;
    margin-bottom: 0.5rem;
    cursor: pointer;
    font-size: 0.85rem;
    color: #94a3b8;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
    transition: all 0.15s;
}
.history-item:hover {
    background: #1a1f2e;
    border-color: #334155;
    color: #e2e8f0;
}
.history-timestamp {
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.65rem;
    color: #475569;
    margin-top: 2px;
}

/* Input box */
.stChatInputContainer {
    background: #13161f !important;
    border: 1px solid #1e2535 !important;
    border-radius: 12px !important;
}

/* Spinner text */
.stSpinner > div {
    border-color: #60a5fa !important;
}

/* Divider */
hr {
    border-color: #1e2535 !important;
    margin: 1rem 0 !important;
}

/* Streamlit chat messages */
[data-testid="stChatMessage"] {
    background: #13161f !important;
    border: 1px solid #1e2535 !important;
    border-radius: 10px !important;
    padding: 0.8rem !important;
    margin-bottom: 0.5rem !important;
}

/* Buttons */
.stButton > button {
    background: #1a1f2e !important;
    color: #94a3b8 !important;
    border: 1px solid #1e2535 !important;
    border-radius: 8px !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 0.78rem !important;
    transition: all 0.15s !important;
}
.stButton > button:hover {
    background: #1e2535 !important;
    color: #e2e8f0 !important;
    border-color: #334155 !important;
}

/* Section headers */
.section-label {
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.7rem;
    letter-spacing: 0.12em;
    color: #475569;
    text-transform: uppercase;
    margin-bottom: 0.8rem;
}

/* Status dots */
.status-dot {
    width: 8px; height: 8px;
    border-radius: 50%;
    display: inline-block;
    margin-right: 6px;
}
.dot-green { background: #4ade80; box-shadow: 0 0 6px #4ade8088; }
.dot-yellow { background: #fbbf24; box-shadow: 0 0 6px #fbbf2488; }
.dot-red { background: #f87171; box-shadow: 0 0 6px #f8717188; }
</style>
""", unsafe_allow_html=True)

# ─── Constants ───────────────────────────────────────────────────────────────
MODELS = {
    "Meta R1": {
        "id": "meta-llama/llama-3.3-70b-instruct",
        "badge_class": "badge-meta",
        "icon": "🧠",
        "desc": "Best for math & reasoning"
    },
    "Qwen3 30B": {
        "id": "qwen/qwen3-next-80b-a3b-instruct",
        "badge_class": "badge-qwen",
        "icon": "⚡",
        "desc": "Strong at science & analysis"
    },
    "Gemma 3 27B": {
        "id": "google/gemma-3-27b-it",
        "badge_class": "badge-gemma",
        "icon": "✨",
        "desc": "Balanced & instruction-tuned"
    }
}

HISTORY_FILE = "chat_history.json"

# ─── History Helpers ─────────────────────────────────────────────────────────
def load_history():
    if os.path.exists(HISTORY_FILE):
        with open(HISTORY_FILE, "r") as f:
            return json.load(f)
    return []

def save_history(history):
    with open(HISTORY_FILE, "w") as f:
        json.dump(history, f, indent=2)

def add_to_history(question, answers: dict):
    history = load_history()
    history.insert(0, {
        "id": str(int(time.time())),
        "question": question,
        "answers": answers,
        "timestamp": datetime.now().strftime("%b %d, %H:%M")
    })
    history = history[:50]   # keep last 50 only
    save_history(history)
    return history

# ─── RAG Loader ──────────────────────────────────────────────────────────────
@st.cache_resource(show_spinner=False)
def load_vectorstore():
    load_dotenv()
    api_key = os.environ.get("OPENROUTER_API_KEY") or os.getenv("OPENROUTER_API_KEY")
    base_url = "https://openrouter.ai/api/v1"
    index_path = "faiss_index"

    if not os.path.exists(index_path):
        return None, "❌ No FAISS index found. Run `rag_engine.py` first."

    embeddings = OpenAIEmbeddings(
        model="qwen/qwen3-embedding-8b",
        openai_api_key=api_key,
        openai_api_base=base_url,
        check_embedding_ctx_length=False
    )
    vectorstore = FAISS.load_local(
        index_path,
        embeddings,
        allow_dangerous_deserialization=True
    )
    return vectorstore, None


def build_chain(vectorstore, model_id: str):
    load_dotenv()
    api_key = os.environ.get("OPENROUTER_API_KEY") or os.getenv("OPENROUTER_API_KEY")
    base_url = "https://openrouter.ai/api/v1"

    llm = ChatOpenAI(
        model=model_id,
        openai_api_key=api_key,
        openai_api_base=base_url,
        temperature=0.3,
        max_tokens=800
    )

    template = """Use the following context to answer the question.
If the answer involves mathematical formulas, format them using LaTeX syntax
($$ for block equations, $ for inline math).
If the answer is NOT in the context, say: "I don't have enough information in the provided documents to answer this."

Context: {context}

Question: {question}

Answer (Markdown + LaTeX if needed):"""

    prompt = PromptTemplate(input_variables=["context", "question"], template=template)

    return RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
        chain_type_kwargs={"prompt": prompt}
    )


def query_model(vectorstore, model_id: str, question: str) -> str:
    try:
        chain = build_chain(vectorstore, model_id)
        result = chain.invoke(question)["result"]
        return result
    except Exception as e:
        if "402" in str(e):
            return "⚠️ Insufficient credits on OpenRouter."
        elif "rate" in str(e).lower():
            return "⚠️ Rate limit hit. Please wait a moment and retry."
        else:
            return f"⚠️ Error: {str(e)}"


# ─── Session State Init ───────────────────────────────────────────────────────
if "chat_history" not in st.session_state:
    st.session_state.chat_history = load_history()

if "current_question" not in st.session_state:
    st.session_state.current_question = None

if "current_answers" not in st.session_state:
    st.session_state.current_answers = {}

if "selected_history" not in st.session_state:
    st.session_state.selected_history = None

# ─── Sidebar ─────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown('<div class="section-label">⚙ System Status</div>', unsafe_allow_html=True)

    vectorstore, vs_error = load_vectorstore()
    if vs_error:
        st.markdown(f'<span class="status-dot dot-red"></span> Index not found', unsafe_allow_html=True)
        st.error(vs_error)
    else:
        st.markdown('<span class="status-dot dot-green"></span> Knowledge base ready', unsafe_allow_html=True)

    st.markdown("---")
    st.markdown('<div class="section-label">🤖 Active Models</div>', unsafe_allow_html=True)
    for name, info in MODELS.items():
        st.markdown(f"{info['icon']} **{name}** — _{info['desc']}_")

    st.markdown("---")
    st.markdown('<div class="section-label">📜 Chat History</div>', unsafe_allow_html=True)

    if st.button("🗑 Clear History", use_container_width=True):
        st.session_state.chat_history = []
        if os.path.exists(HISTORY_FILE):
            os.remove(HISTORY_FILE)
        st.session_state.selected_history = None
        st.rerun()

    st.markdown("")
    for item in st.session_state.chat_history:
        label = item["question"][:42] + "..." if len(item["question"]) > 42 else item["question"]
        if st.button(f"💬 {label}", key=f"hist_{item['id']}", use_container_width=True):
            st.session_state.selected_history = item
            st.session_state.current_question = item["question"]
            st.session_state.current_answers = item["answers"]

# ─── Main Area ───────────────────────────────────────────────────────────────
st.markdown("""
<div class="top-bar">
    <h1>🔬 Multi-Model Research Assistant</h1>
</div>
""", unsafe_allow_html=True)

# Show current Q&A if any
if st.session_state.current_question:
    st.markdown(f"#### 🙋 {st.session_state.current_question}")
    st.markdown("---")

    cols = st.columns(3)
    for i, (name, info) in enumerate(MODELS.items()):
        with cols[i]:
            answer = st.session_state.current_answers.get(name, "_No answer yet._")
            st.markdown(f"""
            <div class="model-card">
                <span class="model-badge {info['badge_class']}">{info['icon']} {name}</span>
                <div class="model-answer">
            """, unsafe_allow_html=True)
            st.markdown(answer)
            st.markdown("</div></div>", unsafe_allow_html=True)

    st.markdown("---")

else:
    # Welcome screen
    st.markdown("""
    <div style="text-align:center; padding: 3rem 0; color: #475569;">
        <div style="font-size: 3rem; margin-bottom: 1rem;">📚</div>
        <div style="font-family: 'Syne', sans-serif; font-size: 1.2rem; color: #64748b;">
            Ask anything — get answers from 3 models simultaneously
        </div>
        <div style="font-size: 0.85rem; margin-top: 0.5rem; color: #334155;">
            Meta R1 &nbsp;·&nbsp; Qwen3 30B &nbsp;·&nbsp; Gemma 3 27B
        </div>
    </div>
    """, unsafe_allow_html=True)

# ─── Chat Input ───────────────────────────────────────────────────────────────
if prompt := st.chat_input("Ask about your documents..."):
    if not vectorstore:
        st.error("Knowledge base not loaded. Please run `rag_engine.py` first.")
    else:
        st.session_state.current_question = prompt
        st.session_state.current_answers = {}
        st.session_state.selected_history = None

        # Show question immediately
        st.markdown(f"#### 🙋 {prompt}")
        st.markdown("---")

        cols = st.columns(3)
        answers = {}

        for i, (name, info) in enumerate(MODELS.items()):
            with cols[i]:
                st.markdown(f"""
                <div class="model-card">
                    <span class="model-badge {info['badge_class']}">{info['icon']} {name}</span>
                """, unsafe_allow_html=True)

                with st.spinner(f"Querying {name}..."):
                    ans = query_model(vectorstore, info["id"], prompt)

                answers[name] = ans
                st.markdown(ans)
                st.markdown("</div>", unsafe_allow_html=True)

        # Save to history
        st.session_state.current_answers = answers
        st.session_state.chat_history = add_to_history(prompt, answers)
        st.rerun()