"""BitNet client — grabbed from https://bitnet-llm-beony.sprites.app"""
import urllib.request
import json

ENDPOINT = "https://bitnet-llm-beony.sprites.app/v1/chat/completions"

def ask(prompt, system=None, max_tokens=256, temperature=0.8):
    """Send a prompt, get a string back."""
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    data = json.dumps({
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
    }).encode()

    req = urllib.request.Request(
        ENDPOINT,
        data=data,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        body = json.loads(resp.read())
    return body["choices"][0]["message"]["content"]


def classify(text, categories):
    """Classify text into one of the given categories."""
    cats = ", ".join(categories)
    return ask(
        f"Classify the following text into exactly one of these categories: {cats}\n\nText: {text}\n\nCategory:",
        system="You are a classifier. Respond with only the category name, nothing else.",
        max_tokens=20,
        temperature=0.1,
    )


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        print(ask(" ".join(sys.argv[1:])))
    else:
        print("Usage: python llm.py \"your prompt here\"")