import json, re, ast, sys, urllib.request

def generate(prompt, model="pycoder"):
    req = urllib.request.Request(
        "http://localhost:11434/api/generate",
        data=json.dumps({"model": model, "prompt": prompt, "stream": False}).encode(),
        headers={"Content-Type": "application/json"},
    )
    text = json.loads(urllib.request.urlopen(req).read())["response"].strip()
    text = re.sub(r"^```(?:python)?\n?|\n?```$", "", text).strip()
    try:
        ast.parse(text)
    except SyntaxError:
        return "# ERROR: Model did not return valid Python."
    return text

if __name__ == "__main__":
    print(generate(" ".join(sys.argv[1:]) or "Write a hello world function"))