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"[KeyManager] Initialized with {len(self.keys)} key(s).")
        if not self.keys:
            print("[KeyManager] 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. Set GROQ_API_KEYS or GROQ_API_KEY in your .env file.")
            key = self.keys[self.current_index]
            self.current_index = (self.current_index + 1) % len(self.keys)
            return key

    def key_count(self) -> int:
        return len(self.keys)


key_manager = KeyRotationManager()