import json
import os
import sys
import urllib.error
import urllib.request

BASE = os.environ["CHAT_BASE"]


class Refused(Exception):
    def __init__(self, status, sentence):
        super().__init__(f"{status}: {sentence}")
        self.status = status


def post(path, body):
    request = urllib.request.Request(
        BASE + path,
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    try:
        return urllib.request.urlopen(request, timeout=60)
    except urllib.error.HTTPError as err:
        raise Refused(err.code, err.read().decode()[:200])


def open_conversation():
    with post("/sessions", {}) as response:
        return json.load(response)["session_id"]


def ask(session_id, message):
    reply, sources, index, ended = "", [], None, None
    with post(f"/sessions/{session_id}/runs", {"message": message}) as stream:
        for line in stream:
            if not line.strip():
                continue
            frame = json.loads(line)
            if frame["type"] == "token":
                reply += frame["delta"]
            elif frame["type"] == "sources":
                sources = [s["name"] for s in frame["sources"]]
            elif frame["type"] == "answer":
                index = frame["index"]
            elif frame["type"] == "run_end":
                ended = frame["end_reason"]
            elif frame["type"] == "error":
                ended = "error"
    return reply, sources, index, ended


if __name__ == "__main__":
    try:
        session_id = open_conversation()
        reply, sources, index, ended = ask(session_id, sys.argv[1])
    except Refused as err:
        sys.exit(str(err))
    print(reply)
    print("sources:", ", ".join(sources))
    if ended != "done":
        sys.exit(f"stopped: {ended}")