"""simple_rag_chatbot_demo.py — Streamlit RAG chatbot over html_notes.txt"""

import subprocess
import sys


def _running_under_streamlit() -> bool:
    try:
        from streamlit.runtime.scriptrunner import get_script_run_ctx
        return get_script_run_ctx() is not None
    except Exception:
        return False


if __name__ == "__main__" and not _running_under_streamlit():
    print("Starting Streamlit server — open http://localhost:8501 in your browser.")
    raise SystemExit(
        subprocess.call(
            [sys.executable, "-m", "streamlit", "run", __file__, *sys.argv[1:]],
        )
    )

from pathlib import Path

import streamlit as st
from dotenv import load_dotenv
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

load_dotenv()

SAMPLES = Path(__file__).parent / "simple_rag_chatbot_samples"
NOTES_PATH = SAMPLES / "html_notes.txt"
DB_DIR = Path(__file__).parent / "simple_rag_chatbot_chroma_db"


def format_docs(docs: list) -> str:
    return "\n\n".join(doc.page_content for doc in docs)


@st.cache_resource
def load_rag_chain():
    if not NOTES_PATH.exists():
        raise FileNotFoundError(
            f"Missing file: {NOTES_PATH}\n"
            "Download simple_rag_chatbot_demo.zip from the lesson page and unzip."
        )

    docs = TextLoader(str(NOTES_PATH), encoding="utf-8").load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=150, chunk_overlap=30)
    chunks = splitter.split_documents(docs)

    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=str(DB_DIR),
    )
    retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

    prompt = ChatPromptTemplate.from_template(
        "Answer in one or two short sentences using only the context below.\n\n"
        "Context:\n{context}\n\n"
        "Question: {question}"
    )
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

    return (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )


st.set_page_config(page_title="HTML Notes Chat")
st.title("HTML Notes Chat")

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

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

if question := st.chat_input("Ask about HTML tags..."):
    st.session_state.messages.append({"role": "human", "content": question})
    st.chat_message("human").write(question)

    chain = load_rag_chain()
    with st.chat_message("ai"):
        answer = chain.invoke(question)
        st.write(answer)

    st.session_state.messages.append({"role": "ai", "content": answer})