import json
import os
import time
from openai import OpenAI

BATCH_IDS_FILE = "data/putnam_eval/submitted_batch_ids.json"

def check_status():
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("Error: OPENAI_API_KEY not set.")
        return
    client = OpenAI(api_key=api_key)

    if not os.path.exists(BATCH_IDS_FILE):
        print(f"Error: {BATCH_IDS_FILE} not found. Did you submit batches?")
        return

    with open(BATCH_IDS_FILE, "r") as f:
        batch_ids = json.load(f)
    
    print(f"Checking status for {len(batch_ids)} batches...\n")
    
    all_completed = True
    completed_count = 0
    
    print(f"{'BATCH ID':<35} | {'STATUS':<15} | {'REQ COUNT':<10} | {'COMPLETED':<10} | {'FAILED':<10}")
    print("-" * 95)
    
    for b_id in batch_ids:
        try:
            batch = client.batches.retrieve(b_id)
            status = batch.status
            counts = batch.request_counts
            total = counts.total if counts else 0
            comp = counts.completed if counts else 0
            fail = counts.failed if counts else 0
            
            print(f"{b_id:<35} | {status:<15} | {total:<10} | {comp:<10} | {fail:<10}")
            
            if status != "completed":
                all_completed = False
            else:
                completed_count += 1
                
        except Exception as e:
            print(f"{b_id:<35} | {'ERROR':<15} | {str(e)}")
            all_completed = False

    print("-" * 95)
    print(f"\nProgress: {completed_count}/{len(batch_ids)} batches completed.")
    
    if all_completed:
        print("\nSUCCESS: All batches finished! You can now run scripts/retrieve_results.py")
    else:
        print("\nSome batches are still processing. Check again later.")

if __name__ == "__main__":
    check_status()