"""tutor.py — Your local RAG tutor with the Kru AI persona baked in.

Combines two techniques from the lesson into one self-contained script:
  - Tech 2 (Modelfile)  — the SYSTEM prompt + parameters are baked into a
                          custom 'kru-ai' Ollama model. Done from Python here,
                          so no separate `ollama create -f Modelfile` step.
  - Tech 3 (LlamaIndex) — RAG over your ./notes folder.

What it does:
  1. On first run, creates the 'kru-ai' model (Modelfile expressed as Python).
  2. Reads every .md file in ./notes and embeds each chunk with
     nomic-embed-text via Ollama.
  3. When you ask a question, it retrieves the closest chunks and asks
     'kru-ai' to answer — so you get RAG _plus_ the Kru AI persona.

Run:
    python tutor.py

Quit with an empty line or Ctrl-C.
"""
from pathlib import Path

import ollama
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.llms.ollama import Ollama

NOTES_DIR = Path("notes")
MODEL_NAME = "kru-ai"
BASE_MODEL = "llama3.2"

# The Modelfile, expressed as Python. The SYSTEM prompt and parameters get
# baked into the 'kru-ai' model on first run — same effect as writing a
# Modelfile and running `ollama create kru-ai -f Modelfile`, but the
# persona lives here as ordinary Python data so this script is self-contained.
SYSTEM_PROMPT = """\
You are Kru AI, a friendly English tutor for 14-year-old students in
Chiang Mai, Thailand.

Rules every reply must follow:
- Answer in simple English. Use vocabulary a Mathayom 2 student knows.
- After your answer, give one important English keyword in [brackets]
  with its Thai translation in (parentheses).
- Keep replies short: 1 to 3 sentences.
- Be warm and encouraging. Use Thai cultural examples when natural
  (food, festivals, temples, school life).
- If the student writes in Thai, reply in English and include the Thai
  version of the key phrase.

Never lecture. Never apologize for being an AI. Just be a kind teacher.
"""
MODEL_PARAMETERS = {"temperature": 0.7, "num_ctx": 4096}


def ensure_kru_ai() -> None:
    """Create the 'kru-ai' model from Python. Idempotent — skips if it
    already exists. EAFP: try to show the model; if that fails, create it."""
    try:
        ollama.show(MODEL_NAME)
        return
    except Exception:
        pass

    print(f"Creating Ollama model '{MODEL_NAME}' (one-time setup)...")
    ollama.create(
        model=MODEL_NAME,
        from_=BASE_MODEL,
        system=SYSTEM_PROMPT,
        parameters=MODEL_PARAMETERS,
    )
    print(f"  done. Run `ollama list` to see {MODEL_NAME}.")


def main() -> None:
    ensure_kru_ai()

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

    if not NOTES_DIR.exists() or not any(NOTES_DIR.glob("*.md")):
        print(f"No .md files found in {NOTES_DIR}/")
        print("Create some markdown notes there first, then re-run.")
        return

    print(f"Reading notes from {NOTES_DIR}/ ...")
    docs = SimpleDirectoryReader(
        str(NOTES_DIR), required_exts=[".md"], recursive=True
    ).load_data()
    print(f"Loaded {len(docs)} document(s). Building index — please wait ~30s.")

    index = VectorStoreIndex.from_documents(docs)
    engine = index.as_query_engine(similarity_top_k=3)
    print("Ready. Ask me anything about your notes. Empty line to quit.\n")

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

        response = engine.query(question)
        print(f"\nbot > {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()