import os
import threading
from typing import List

class KeyRotationManager:
    def __init__(self):
        keys_str = os.getenv("GROQ_API_KEYS")
        if not keys_str:
            keys_str = os.getenv("GROQ_API_KEY", "")
        
        self.keys = [k.strip() for k in keys_str.split(",") if k.strip()]
        print(f"KeyRotationManager initialized with {len(self.keys)} keys.")
        if not self.keys:
            print("WARNING: No GROQ_API_KEYS or GROQ_API_KEY found in environment!")
        self.current_index = 0
        self.lock = threading.Lock()

    def get_next_key(self) -> str:
        with self.lock:
            if not self.keys:
                raise ValueError("No GROQ_API_KEYS found in environment variables.")
            
            key = self.keys[self.current_index]
            self.current_index = (self.current_index + 1) % len(self.keys)
            return key

key_manager = KeyRotationManager()