import argparse
import os
import time

from openai import OpenAI

parser = argparse.ArgumentParser(description="Progress")

parser.add_argument("--batch_id", type=str, required=True, help="Input file path")

args = parser.parse_args()


os.environ["OPENAI_API_KEY"] = (
    "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"
)

client = OpenAI()

with open(f"batchapi/jobs/jobs_1/{args.batch_id}", "r") as f:
    # batch_id = "batch_69899f6279a48190abaad4d65407768b"
    batch_id = f.readline().strip()

print(batch_id)

while True:
    b = client.batches.retrieve(batch_id)
    print(f"Status: {b.status}, {b.request_counts}", end="\r", flush=True)
    # print("Status:", b.status)

    if b.status in ["completed", "failed", "cancelled"]:
        break

    time.sleep(1)
print("Status:", b.status)