"""karpathy_tutor.py — Same idea as tutor.py, but the corpus is the
blog posts from karpathy.github.io.

Before running:
    git clone https://github.com/karpathy/karpathy.github.io karpathy_blog

Then:
    python karpathy_tutor.py

The first run takes ~3–5 minutes because every blog post has to be
embedded. The index is saved to ./karpathy_index/, so the second run
loads in about a second.

Delete the karpathy_index/ folder if you ever change the source files
and want to re-index from scratch.
"""
from pathlib import Path

from llama_index.core import (
    Settings,
    SimpleDirectoryReader,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.llms.ollama import Ollama

BLOG_DIR = Path("karpathy_blog/_posts")
PERSIST_DIR = Path("karpathy_index")

Settings.llm = Ollama(model="llama3.2", request_timeout=180)
Settings.embed_model = OllamaEmbedding(model_name="nomic-embed-text")


def build_or_load_index() -> VectorStoreIndex:
    if PERSIST_DIR.exists() and any(PERSIST_DIR.iterdir()):
        print(f"Loading saved index from {PERSIST_DIR}/ ...")
        ctx = StorageContext.from_defaults(persist_dir=str(PERSIST_DIR))
        return load_index_from_storage(ctx)

    if not BLOG_DIR.exists():
        raise SystemExit(
            f"Missing {BLOG_DIR}. First run:\n"
            "  git clone https://github.com/karpathy/karpathy.github.io karpathy_blog"
        )

    print(f"Reading blog posts from {BLOG_DIR}/ ...")
    docs = SimpleDirectoryReader(
        str(BLOG_DIR), required_exts=[".md", ".markdown"], recursive=True
    ).load_data()
    print(f"Loaded {len(docs)} post(s). Embedding — this can take 3-5 minutes.")
    index = VectorStoreIndex.from_documents(docs)
    PERSIST_DIR.mkdir(exist_ok=True)
    index.storage_context.persist(str(PERSIST_DIR))
    print(f"Saved index to {PERSIST_DIR}/ — next run will be instant.\n")
    return index


def main() -> None:
    index = build_or_load_index()
    engine = index.as_query_engine(similarity_top_k=3)
    print("Ready. Try: What is Software 2.0?  /  Why should I understand backprop?\n")

    while True:
        try:
            question = input("you > ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return
        if not question:
            return

        response = engine.query(question)
        print(f"\nKarpathy-bot > {response}\n")
        sources = {
            (n.metadata or {}).get("file_name", "?") for n in response.source_nodes
        }
        if sources:
            print("              sources:", ", ".join(sorted(sources)), "\n")


if __name__ == "__main__":
    main()