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

parser = argparse.ArgumentParser(description="Stream a reply from local Ollama.")
parser.add_argument("prompt")
args = parser.parse_args()

payload = {
    "model": "qwen2.5:1.5b",
    "messages": [{"role": "user", "content": args.prompt}],
    "options": {"num_predict": 160},
    "stream": True,
}
request = Request(
    "http://127.0.0.1:11434/api/chat",
    data=json.dumps(payload).encode("utf-8"),
    headers={"Content-Type": "application/json"},
    method="POST",
)

started = perf_counter()
first_text_at = None
finished = False
try:
    with urlopen(request, timeout=120) as response:
        for line in response:
            if not line.strip():
                continue
            event = json.loads(line)
            if event.get("error"):
                raise SystemExit(f"\nOllama stopped the stream: {event['error']}")
            text = event.get("message", {}).get("content", "")
            if text:
                if first_text_at is None:
                    first_text_at = perf_counter() - started
                print(text, end="", flush=True)
            if event.get("done"):
                finished = True
                break
except HTTPError as error:
    raise SystemExit(f"Ollama returned HTTP {error.code}. Check the model tag.") from None
except (URLError, OSError, ValueError) as error:
    raise SystemExit(f"\nCould not finish the stream: {error}") from None

print()
if not finished:
    raise SystemExit("The connection closed before Ollama reported completion.")
if first_text_at is None:
    raise SystemExit("Ollama finished without emitting visible text.")
print(f"First text in {first_text_at:.2f} s.", file=sys.stderr)