"""ask.py — talk to your own AI, running on your own machine.

The AI is Kru Fah, an English tutor. She always answers in English, even when
a student writes to her in Thai. That keeps every answer inside the lesson.

The role does NOT live in this file. It lives in the Modelfile, which Ollama
reads when it builds the model. Look: this program never sends a system
prompt. The teacher personality still comes out, because it is baked into the
model itself.

Set it up once:

    ollama pull qwen2.5:3b
    ollama create krueng-tutor -f Modelfile

Then:

    python3 ask.py "What is the past perfect?"
    python3 ask.py --chat                    # keep talking
    python3 ask.py --who                     # prove the role came from the Modelfile
    python3 ask.py --host http://spark.local:11434 "hello"

Nothing you type here goes to the internet.
"""
import argparse
import sys

from llama_index.core.llms import ChatMessage
from llama_index.llms.ollama import Ollama

# The tutor answers in English, but students type Thai, and a model can always
# surprise you. On Windows one Thai character stops the program dead with a
# "charmap codec" error. This line fixes that. Keep it.
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

DEFAULT_MODEL = "krueng-tutor"
DEFAULT_HOST = "http://localhost:11434"


def build_llm(model, host):
    """Make the LlamaIndex object that speaks to Ollama.

    Notice what is NOT here: system_prompt. We leave it out on purpose.
    Ollama then uses the SYSTEM block from the Modelfile.
    """
    return Ollama(
        model=model,
        base_url=host,
        request_timeout=180.0,   # a big model on a cold start is slow
    )


def say(llm, history):
    """Send the conversation and print the answer as it arrives.

    We ask the model in the Modelfile not to use markdown stars. It keeps
    using them anyway. So we stop asking and just take them out here.
    A rule you can enforce in code beats a rule you have to hope for.
    """
    answer = ""
    for chunk in llm.stream_chat(history):
        piece = (chunk.delta or "").replace("*", "")
        print(piece, end="", flush=True)   # flush, or nothing shows until the end
        answer += piece
    print()
    return answer


def main():
    parser = argparse.ArgumentParser(description="Talk to your own local AI.")
    parser.add_argument("question", nargs="*", help="what to ask")
    parser.add_argument("--chat", action="store_true", help="keep the conversation going")
    parser.add_argument("--who", action="store_true", help="ask the model who it is")
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--host", default=DEFAULT_HOST)
    args = parser.parse_args()

    llm = build_llm(args.model, args.host)
    history = []

    try:
        if args.who:
            history.append(ChatMessage(role="user", content="Who are you? Answer in one sentence."))
            say(llm, history)
            return 0

        if args.question:
            history.append(ChatMessage(role="user", content=" ".join(args.question)))
            answer = say(llm, history)
            history.append(ChatMessage(role="assistant", content=answer))

        if not args.chat:
            if not args.question:
                parser.print_help()
            return 0

        print("\nType a question. Press Ctrl+C to stop.\n")
        while True:
            try:
                question = input("you  > ").strip()
            except (EOFError, KeyboardInterrupt):
                print("\nBye!")
                return 0
            if not question:
                continue
            history.append(ChatMessage(role="user", content=question))
            print("kru  > ", end="", flush=True)
            answer = say(llm, history)
            history.append(ChatMessage(role="assistant", content=answer))
            print()

    except Exception as problem:
        return complain(problem, args)


def complain(problem, args):
    """Turn an ugly traceback into one useful sentence."""
    text = str(problem).lower()
    print()
    if "connect" in text or "refused" in text or "timed out" in text:
        print("Cannot reach Ollama at %s" % args.host)
        print("Is it running? Try:  ollama list")
        print("Reaching another computer? That machine needs OLLAMA_HOST=0.0.0.0")
    elif "not found" in text or "no such" in text:
        print("Ollama does not have a model called '%s'." % args.model)
        print("Build it:  ollama create %s -f Modelfile" % args.model)
    else:
        print("Something went wrong: %s" % problem)
    return 1


if __name__ == "__main__":
    sys.exit(main())