"""end_to_end_rag_demo.py — loader → splitter → embeddings → Chroma → retriever → chain"""

import shutil
from pathlib import Path

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 / "end_to_end_rag_samples"
NOTES_PATH = SAMPLES / "html_notes.txt"
QUESTION = "How do I make a link on a page?"
DB_DIR = Path(__file__).parent / "end_to_end_rag_chroma_db"
CHUNK_SIZE = 150
CHUNK_OVERLAP = 30


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


def main() -> None:
    if not NOTES_PATH.exists():
        raise SystemExit(
            f"Missing file: {NOTES_PATH}\n"
            "Download end_to_end_rag_demo.zip from the lesson page and unzip."
        )

    if DB_DIR.exists():
        shutil.rmtree(DB_DIR)

    docs = TextLoader(str(NOTES_PATH), encoding="utf-8").load()
    print(f"=== Loader ===\n1 document, {len(docs[0].page_content)} characters")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )
    chunks = splitter.split_documents(docs)
    print(f"\n=== Text Splitter ===\n{len(chunks)} chunks")

    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=str(DB_DIR),
    )
    print(f"\n=== Vector Store ===\nSaved to {DB_DIR}")

    retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

    print("\n=== Retriever ===")
    retrieved = retriever.invoke(QUESTION)
    for i, doc in enumerate(retrieved):
        text = doc.page_content.replace("\n", " ").strip()
        snippet = text[:80] + ("…" if len(text) > 80 else "")
        print(f"[{i}] {snippet}")

    prompt = ChatPromptTemplate.from_template(
        "Answer in one short sentence using only the context below.\n\n"
        "Context:\n{context}\n\n"
        "Question: {question}"
    )

    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

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

    print("\n=== Chain output ===")
    print(f"Q: {QUESTION}")
    result = rag_chain.invoke(QUESTION)
    print(f"Printed: {result}")


if __name__ == "__main__":
    main()