import argparse
import json
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

BASE = "http://127.0.0.1:11434/api"

def request(endpoint, payload):
    req = Request(BASE + endpoint, data=json.dumps(payload).encode("utf-8"),
                  headers={"Content-Type": "application/json"}, method="POST")
    with urlopen(req, timeout=180) as response:
        return json.load(response)

def main():
    parser = argparse.ArgumentParser(description="Chat with an installed local Ollama model")
    parser.add_argument("--model", required=True, help="Exact installed local model name")
    args = parser.parse_args()
    if "cloud" in args.model.lower():
        parser.exit(2, "Choose a downloaded local model, not a cloud model.\n")
    messages = [{"role": "system", "content": "You are a study helper. Say when you are unsure."}]
    try:
        info = request("/show", {"model": args.model})
        if info.get("remote_host") or info.get("remote_model"):
            parser.exit(2, "Remote models are not supported in this local tutorial.\n")
        print("Type /quit to leave. Recent conversation is kept only in memory.")
        while True:
            text = input("You: ").strip()
            if text == "/quit":
                break
            if not text:
                continue
            messages = [messages[0]] + messages[1:][-12:]
            messages.append({"role": "user", "content": text})
            result = request("/chat", {"model": args.model, "messages": messages, "stream": False})
            reply = result.get("message", {}).get("content")
            if not isinstance(reply, str):
                raise ValueError("The local server returned no text reply.")
            print("Assistant:", reply)
            messages.append({"role": "assistant", "content": reply})
    except (HTTPError, URLError, TimeoutError, OSError, ValueError):
        parser.exit(1, "Local model request failed. Check Ollama, the installed model name, and available memory.\n")
    except (EOFError, KeyboardInterrupt):
        print("\nChat closed.")

if __name__ == "__main__":
    main()