import threading
import time
import asyncio
import logging
import os

logger = logging.getLogger(__name__)

class ThreadSafeTokenManager:
    def __init__(self, token_sets=None, use_database=False):
        """
        Initialize token manager.
        If use_database=True, tokens will be loaded from Laravel database.
        Otherwise, uses the token_sets list from config.py
        """
        self.use_database = use_database
        self._db_loader = None
        self._token_sets = token_sets or []

        if self.use_database:
            try:
                from app.db_token_loader import DatabaseTokenLoader
                self._db_loader = DatabaseTokenLoader(
                    api_base_url=os.environ.get('LARAVEL_API_URL', 'http://127.0.0.1:9000'),
                    api_key=os.environ.get('ADMIN_API_KEY', 'locket-gold-secret-2024')
                )
                logger.info("[TokenManager] Using database token loader")
            except ImportError:
                logger.warning("[TokenManager] Could not import DatabaseTokenLoader, using config fallback")
                self.use_database = False

        # Initialize token state tracking
        self.last_used = {i: 0.0 for i in range(len(self._token_sets))}
        self.locked_tokens = set()  # Set of token indices currently in use
        self.lock = threading.Lock()  # Global thread lock for state updates
        self.current_idx = 0
        self._reload_interval = 300  # Reload from DB every 5 minutes
        self._last_db_reload = 0

    async def _ensure_tokens_loaded(self):
        """Ensure tokens are loaded from database if using DB mode"""
        if not self.use_database or not self._db_loader:
            return

        now = time.time()
        if now - self._last_db_reload > self._reload_interval or not self._token_sets:
            try:
                db_tokens = await self._db_loader.fetch_tokens()
                if db_tokens:
                    # Rebuild token list from database
                    old_count = len(self._token_sets)
                    self._token_sets = db_tokens

                    # Rebuild last_used tracking for new tokens
                    new_last_used = {}
                    for i, token in enumerate(self._token_sets):
                        new_last_used[i] = self.last_used.get(i, 0.0)

                    # Add entries for new tokens
                    for i in range(len(self._token_sets)):
                        if i not in new_last_used:
                            new_last_used[i] = 0.0

                    self.last_used = new_last_used
                    self._last_db_reload = now

                    # Remove locked status for tokens no longer in DB
                    self.locked_tokens = {i for i in self.locked_tokens if i < len(self._token_sets)}

                    logger.info(f"[TokenManager] Reloaded tokens from DB: {old_count} -> {len(self._token_sets)}")
            except Exception as e:
                logger.error(f"[TokenManager] Failed to reload tokens from DB: {e}")

    async def acquire_token(self, cooldown_seconds=15):
        """
        Acquires an available token in a thread-safe and loop-safe way.
        If all tokens are busy, it waits in an asyncio-friendly way and retries.
        """
        # Load tokens from database if needed
        await self._ensure_tokens_loaded()

        wait_to_use = 0.0
        idx_to_use = None
        token_to_use = None

        while True:
            with self.lock:
                n = len(self._token_sets)
                if n == 0:
                    logger.error("[TokenManager] No tokens available!")
                    raise ValueError("No tokens configured in TOKEN_SETS!")

                now = time.time()
                best_idx = None
                best_wait = 999999.0

                # Check tokens in a round-robin starting from current_idx
                for offset in range(n):
                    idx = (self.current_idx + offset) % n
                    if idx not in self.locked_tokens:
                        # Get token-specific cooldown or use default
                        token = self._token_sets[idx]
                        token_cooldown = token.get('cooldown_seconds', cooldown_seconds) if isinstance(token, dict) else cooldown_seconds

                        # Calculate cooldown remaining
                        elapsed = now - self.last_used[idx]
                        wait_time = max(0.0, token_cooldown - elapsed)

                        if wait_time < best_wait:
                            best_wait = wait_time
                            best_idx = idx

                if best_idx is not None:
                    # Mark the token as currently locked
                    self.locked_tokens.add(best_idx)
                    self.current_idx = (best_idx + 1) % n

                    idx_to_use = best_idx
                    token_to_use = self._token_sets[best_idx]
                    wait_to_use = best_wait
                    break

            # If all tokens are busy (currently executing another request),
            # sleep for 1 second asynchronously and retry.
            await asyncio.sleep(1.0)

        # Handle the cooldown sleep in an asyncio-friendly manner.
        # This yields control so other tasks/requests can proceed.
        if wait_to_use > 0:
            logger.info(f"[TokenManager] Cooldown active for Token #{idx_to_use + 1}. Waiting {wait_to_use:.2f}s...")
            await asyncio.sleep(wait_to_use)

        with self.lock:
            # Update the last used timestamp when the request actually begins
            self.last_used[idx_to_use] = time.time()

        # Mark token as used in database if using DB mode
        if self.use_database and self._db_loader and isinstance(token_to_use, dict):
            token_id = token_to_use.get('id')
            if token_id and token_id > 0:  # Only for DB tokens (positive IDs)
                asyncio.create_task(self._db_loader.mark_token_used(token_id))

        logger.info(f"[TokenManager] Acquired Token #{idx_to_use + 1}")
        return idx_to_use, token_to_use

    def release_token(self, token_idx):
        """
        Releases the lock on the specified token index.
        """
        with self.lock:
            if token_idx in self.locked_tokens:
                self.locked_tokens.remove(token_idx)
                logger.info(f"[TokenManager] Released Token #{token_idx + 1}")

    def reload_from_database(self):
        """Force reload tokens from database"""
        if self.use_database:
            self._last_db_reload = 0  # Will trigger reload on next acquire
            logger.info("[TokenManager] Database reload scheduled")