import os

from colabhive import ColabHive

client = ColabHive(
    os.environ["COLABHIVE_API_KEY"],
    os.environ["COLABHIVE_ACCOUNT_ID"],
    base_url=os.getenv("COLABHIVE_API_URL", "https://api.colabhive.com"),
)
endpoint_id = os.environ["COLABHIVE_ENDPOINT_ID"]
capability = client.cohorts.capabilities(endpoint_id)
if not capability.available:
    raise RuntimeError(f"Cohort unavailable: {capability.reason}")

accepted = client.cohorts.run(
    endpoint_id=endpoint_id,
    input_data={"text": "hello"},
    execution={"mode": "cohort", "fallback": "single", "trace": "summary"},
    idempotency_key="cohort-example-0001",
)
print(client.cohorts.wait(accepted.task_id))