import json
import os
import time
from pathlib import Path

import requests

base = os.getenv("LAB0_BASE_URL", "https://app.lab0.ai/app").rstrip("/")
namespace = os.getenv("LAB0_NAMESPACE", "quickstart")
path = Path(os.getenv("LAB0_FILE_PATH", "plan.md"))
session = requests.Session()
session.headers["Authorization"] = f"Bearer {os.environ['LAB0_API_KEY']}"


def request(method, route, **kwargs):
    response = session.request(method, base + route, timeout=200, **kwargs)
    if not response.ok:
        raise RuntimeError(f"HTTP {response.status_code}: {response.text}")
    return response.json()


with path.open("rb") as document:
    uploaded = request("POST", "/api/knowledge/files",
                       files={"file": (path.name, document)},
                       data={"namespace": namespace})
file_id = uploaded["id"]
for _ in range(90):
    document = request("GET", f"/api/knowledge/files/{file_id}")
    if document["status"] == "ready":
        break
    if document["status"] == "failed":
        raise RuntimeError(document.get("error", "Processing failed"))
    time.sleep(2)
else:
    raise TimeoutError(f"Processing timed out; inspect file {file_id}")

answer = request("POST", "/api/knowledge/respond", json={
    "prompt": "How is the rollout run?",
    "namespace": namespace,
    "file_ids": [file_id],
})
print(json.dumps(answer, indent=2))