import os
import logging
from typing import List, Optional, Dict
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class InstanceProvider:
    """Manages multiple Cerebras API keys with simple rotation"""
    
    def __init__(self):
        self.api_keys: List[str] = []
        self.current_index = 0
        self.base_url = os.getenv("CEREBRAS_BASE_URL")
        self.model_name = os.getenv("CEREBRAS_MODEL", "llama3.1-70b")
        self._initialize_instances()
    
    def _initialize_instances(self):
        """Load all API keys into a list"""
        keys_str = os.getenv("CEREBRAS_API_KEYS", "")
        self.api_keys = [k.strip() for k in keys_str.split(",") if k.strip()]
        
        if not self.api_keys:
            logger.error("No API keys found in CEREBRAS_API_KEYS")

    def get_next_instance(self) -> Optional[Dict[str, str]]:
        """
        Returns a dictionary with the credentials for the next instance.
        Returns: {'api_key': str, 'base_url': str, 'model': str}
        """
        if not self.api_keys:
            return None
        
        # Get current key
        key = self.api_keys[self.current_index]
        
        # Rotate index for the next call
        self.current_index = (self.current_index + 1) % len(self.api_keys)
        
        return {
            "api_key": key,
            "base_url": self.base_url,
            "model": self.model_name
        }
    
    def get_total_instances(self) -> int:
        return len(self.api_keys)