"""
Adaptive rate limit tracker that learns optimal retry wait times for each search engine.
"""

import random
import threading
import time
from collections import deque
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

from cachetools import LRUCache

if TYPE_CHECKING:
    from sqlalchemy.orm import Session
from loguru import logger

from ...config.thread_settings import (
    NoSettingsContextError,
    get_setting_from_snapshot,
    get_settings_context,
)
from ...settings.env_registry import is_ci_environment
from ...utilities.thread_context import get_search_context

# Lazy imports to avoid database initialization in programmatic mode
_db_imports = None


def _get_db_imports():
    """Lazy load database imports only when needed."""
    global _db_imports
    if _db_imports is None:
        try:
            from ...database.models import RateLimitAttempt, RateLimitEstimate
            from ...database.session_context import get_user_db_session

            _db_imports = {
                "RateLimitAttempt": RateLimitAttempt,
                "RateLimitEstimate": RateLimitEstimate,
                "get_user_db_session": get_user_db_session,
            }
        except (ImportError, RuntimeError):
            # Database not available - programmatic mode
            _db_imports = {}
    return _db_imports


class AdaptiveRateLimitTracker:
    """
    Tracks and learns optimal retry wait times for each search engine.
    Persists learned patterns to the main application database using SQLAlchemy.
    """

    def __init__(self, settings_snapshot=None, programmatic_mode=False):
        self.settings_snapshot = settings_snapshot or {}
        self.programmatic_mode = programmatic_mode

        # Helper function to get settings with defaults
        def get_setting_or_default(key, default, type_fn=None):
            try:
                value = get_setting_from_snapshot(
                    key,
                    settings_snapshot=self.settings_snapshot,
                )
                # Handle None values - return default instead of calling type_fn(None)
                if value is None:
                    return default
                return type_fn(value) if type_fn else value
            except NoSettingsContextError:
                return default

        # Get settings with explicit defaults
        self.memory_window = get_setting_or_default(
            "rate_limiting.memory_window", 100, int
        )
        self.exploration_rate = get_setting_or_default(
            "rate_limiting.exploration_rate", 0.1, float
        )
        self.learning_rate = get_setting_or_default(
            "rate_limiting.learning_rate", 0.3, float
        )
        self.decay_per_day = get_setting_or_default(
            "rate_limiting.decay_per_day", 0.95, float
        )

        # In programmatic mode, default to disabled
        self.enabled = get_setting_or_default(
            "rate_limiting.enabled",
            not self.programmatic_mode,  # Default based on mode
            bool,
        )

        profile = get_setting_or_default("rate_limiting.profile", "balanced")

        if self.programmatic_mode and self.enabled:
            logger.info(
                "Rate limiting enabled in programmatic mode - using memory-only tracking without persistence"
            )

        # Apply rate limiting profile
        self._apply_profile(profile)

        # Bounded in-memory caches for fast access (prevent memory leaks)
        # Lock required: LRUCache is not thread-safe (mutates on read via LRU reordering)
        self._cache_lock = threading.Lock()
        self.recent_attempts: LRUCache = LRUCache(maxsize=100)
        self.current_estimates: LRUCache = LRUCache(maxsize=100)

        # Initialize the _estimates_loaded flag
        self._estimates_loaded = False

        # Load estimates from database
        self._load_estimates()

        logger.info(
            f"AdaptiveRateLimitTracker initialized: enabled={self.enabled}, profile={profile}"
        )

    def _apply_profile(self, profile: str) -> None:
        """Apply rate limiting profile settings."""
        if profile == "conservative":
            # More conservative: lower exploration, slower learning
            self.exploration_rate = min(
                self.exploration_rate * 0.5, 0.05
            )  # 5% max exploration
            self.learning_rate = min(
                self.learning_rate * 0.7, 0.2
            )  # Slower learning
            logger.info("Applied conservative rate limiting profile")
        elif profile == "aggressive":
            # More aggressive: higher exploration, faster learning
            self.exploration_rate = min(
                self.exploration_rate * 1.5, 0.2
            )  # Up to 20% exploration
            self.learning_rate = min(
                self.learning_rate * 1.3, 0.5
            )  # Faster learning
            logger.info("Applied aggressive rate limiting profile")
        else:  # balanced
            # Use settings as-is
            logger.info("Applied balanced rate limiting profile")

    def _load_estimates(self) -> None:
        """Load estimates from database into memory."""
        # Skip database operations in programmatic mode
        if self.programmatic_mode:
            logger.debug(
                "Skipping rate limit estimate loading in programmatic mode"
            )
            self._estimates_loaded = (
                True  # Mark as loaded to skip future attempts
            )
            return

        # During initialization, we don't have user context yet
        # Mark that we need to load estimates when user context becomes available
        self._estimates_loaded = False
        logger.debug(
            "Rate limit estimates will be loaded on-demand when user context is available"
        )

    def _ensure_estimates_loaded(self) -> None:
        """Load estimates from user's encrypted database if not already loaded."""
        # Early return if already loaded or should skip
        if self._estimates_loaded or self.programmatic_mode:
            if not self._estimates_loaded:
                self._estimates_loaded = True
            return

        # Get database imports
        db_imports = _get_db_imports()
        RateLimitEstimate = (
            db_imports.get("RateLimitEstimate") if db_imports else None
        )

        if not db_imports or not RateLimitEstimate:
            # Database not available
            self._estimates_loaded = True
            return

        # Try to get research context from search tracker

        context = get_search_context()
        if not context:
            # No context available (e.g., in tests or programmatic access)
            self._estimates_loaded = True
            return

        username = context.get("username")
        password = context.get("user_password")

        if username and password:
            try:
                # Use thread-safe metrics writer to read from user's encrypted database
                from ...database.thread_metrics import metrics_writer

                # Set password for this thread
                metrics_writer.set_user_password(username, password)

                session: "Session"
                with metrics_writer.get_session(username) as session:
                    estimates: list[Any] = session.query(
                        RateLimitEstimate
                    ).all()

                    for estimate in estimates:
                        # Apply decay for old estimates
                        age_hours = (time.time() - estimate.last_updated) / 3600
                        decay = self.decay_per_day ** (age_hours / 24)

                        with self._cache_lock:
                            self.current_estimates[estimate.engine_type] = {
                                "base": estimate.base_wait_seconds,
                                "min": estimate.min_wait_seconds,
                                "max": estimate.max_wait_seconds,
                                "confidence": decay,
                            }

                        logger.debug(
                            f"Loaded estimate for {estimate.engine_type}: base={estimate.base_wait_seconds:.2f}s, confidence={decay:.2f}"
                        )

                self._estimates_loaded = True
                logger.info(
                    f"Loaded {len(estimates)} rate limit estimates from encrypted database"
                )

            except Exception:
                logger.warning("Could not load rate limit estimates")
                # Mark as loaded anyway to avoid repeated attempts
                self._estimates_loaded = True

    def get_wait_time(self, engine_type: str) -> float:
        """
        Get adaptive wait time for a search engine.
        Includes exploration to discover better rates.

        Args:
            engine_type: Name of the search engine

        Returns:
            Wait time in seconds
        """
        # If rate limiting is disabled, return minimal wait time
        if not self.enabled:
            return 0.1

        # Check if we have a user context - if not, handle appropriately
        context = get_search_context()
        if not context and not self.programmatic_mode:
            # No context and not in programmatic mode - this is unexpected
            logger.warning(
                f"No user context available for rate limiting on {engine_type} "
                "but programmatic_mode=False. Disabling rate limiting. "
                "This may indicate a configuration issue."
            )
            return 0.0

        # In programmatic mode, we continue with memory-only rate limiting even without context

        # Ensure estimates are loaded from database
        self._ensure_estimates_loaded()

        with self._cache_lock:
            if engine_type not in self.current_estimates:
                # First time seeing this engine - start optimistic and learn from real responses
                # Use engine-specific optimistic defaults only for what we know for sure
                optimistic_defaults = {
                    "SearXNGSearchEngine": 0.1,  # Self-hosted default engine
                }

                wait_time = optimistic_defaults.get(
                    engine_type, 0.1
                )  # Default optimistic for others
                logger.info(
                    f"No rate limit data for {engine_type}, starting optimistic with {wait_time}s"
                )
                return wait_time

            estimate = self.current_estimates[engine_type]
            base_wait = estimate["base"]

        # Security: random used for non-security rate-limit jitter, not tokens or secrets
        # Exploration vs exploitation
        if random.random() < self.exploration_rate:
            # Explore: try a faster rate to see if API limits have relaxed
            wait_time = base_wait * random.uniform(0.5, 0.9)
            logger.debug(
                f"Exploring faster rate for {engine_type}: {wait_time:.2f}s"
            )
        else:
            # Exploit: use learned estimate with jitter
            wait_time = base_wait * random.uniform(0.9, 1.1)

        # Enforce bounds
        wait_time = max(
            float(estimate["min"]), min(wait_time, float(estimate["max"]))
        )
        return float(wait_time)

    def apply_rate_limit(self, engine_type: str) -> float:
        """
        Apply rate limiting for the given engine type.
        This is a convenience method that combines checking if rate limiting
        is enabled, getting the wait time, and sleeping if necessary.

        Args:
            engine_type: The type of search engine

        Returns:
            The wait time that was applied (0 if rate limiting is disabled)
        """
        if not self.enabled:
            return 0.0

        wait_time = self.get_wait_time(engine_type)
        if wait_time > 0:
            logger.debug(
                f"{engine_type} waiting {wait_time:.2f}s before request"
            )
            time.sleep(wait_time)
        return wait_time

    def record_outcome(
        self,
        engine_type: str,
        wait_time: float,
        success: bool,
        retry_count: int,
        error_type: Optional[str] = None,
        search_result_count: Optional[int] = None,
    ) -> None:
        """
        Record the outcome of a retry attempt.

        Args:
            engine_type: Name of the search engine
            wait_time: How long we waited before this attempt
            success: Whether the attempt succeeded
            retry_count: Which retry attempt this was (1, 2, 3, etc.)
            error_type: Type of error if failed
            search_result_count: Number of search results returned (for quality monitoring)
        """
        # If rate limiting is disabled, don't record outcomes
        if not self.enabled:
            logger.info(
                f"Rate limiting disabled - not recording outcome for {engine_type}"
            )
            return

        logger.debug(
            f"Recording rate limit outcome for {engine_type}: success={success}, wait_time={wait_time}s"
        )
        timestamp = time.time()

        # NOTE: Database writes for rate limiting are disabled to prevent
        # database locking issues under heavy parallel search load.
        # Rate limiting still works via in-memory tracking below.
        # Historical rate limit data is not persisted to DB.

        # Update in-memory tracking
        with self._cache_lock:
            if engine_type not in self.recent_attempts:
                # Get current memory window setting from thread context
                settings_context = get_settings_context()
                if settings_context:
                    current_memory_window = int(
                        settings_context.get_setting(
                            "rate_limiting.memory_window", self.memory_window
                        )
                    )
                else:
                    current_memory_window = self.memory_window

                self.recent_attempts[engine_type] = deque(
                    maxlen=current_memory_window
                )

            self.recent_attempts[engine_type].append(
                {
                    "wait_time": wait_time,
                    "success": success,
                    "timestamp": timestamp,
                    "retry_count": retry_count,
                    "search_result_count": search_result_count,
                }
            )

        # Update estimates
        self._update_estimate(engine_type)

    def _update_estimate(self, engine_type: str) -> None:
        """Update wait time estimate based on recent attempts."""
        with self._cache_lock:
            if (
                engine_type not in self.recent_attempts
                or len(self.recent_attempts[engine_type]) < 3
            ):
                attempt_count = len(self.recent_attempts.get(engine_type) or [])
                logger.info(
                    f"Not updating estimate for {engine_type} - only {attempt_count} attempts (need 3)"
                )
                return

            attempts = list(self.recent_attempts[engine_type])
            old_estimate = self.current_estimates.get(engine_type)

        # Calculate success rate and optimal wait time
        successful_waits = [a["wait_time"] for a in attempts if a["success"]]
        failed_waits = [a["wait_time"] for a in attempts if not a["success"]]

        if not successful_waits:
            # All attempts failed - increase wait time with a cap
            new_base = max(failed_waits) * 1.5 if failed_waits else 10.0
            # Cap the base wait time to prevent runaway growth
            new_base = min(new_base, 10.0)  # Max 10 seconds base when all fail
        else:
            # Use 50th percentile (median) of successful waits for more stability
            # This provides a balanced approach between speed and reliability
            successful_waits.sort()
            percentile_50 = successful_waits[
                max(0, int(len(successful_waits) * 0.50) - 1)
            ]
            new_base = percentile_50

        # Update estimate with learning rate (exponential moving average)
        if old_estimate is not None:
            old_base = old_estimate["base"]
            # Get current learning rate from settings context
            settings_context = get_settings_context()
            if settings_context:
                current_learning_rate = float(
                    settings_context.get_setting(
                        "rate_limiting.learning_rate", self.learning_rate
                    )
                )
            else:
                current_learning_rate = self.learning_rate

            new_base = (
                1 - current_learning_rate
            ) * old_base + current_learning_rate * new_base

        # Apply absolute cap to prevent extreme wait times
        new_base = min(new_base, 10.0)  # Cap base at 10 seconds

        # Calculate bounds with more reasonable limits
        min_wait = max(0.01, new_base * 0.5)
        max_wait = min(10.0, new_base * 3.0)  # Max 10 seconds absolute cap

        # Update in memory
        with self._cache_lock:
            self.current_estimates[engine_type] = {
                "base": new_base,
                "min": min_wait,
                "max": max_wait,
                "confidence": min(len(attempts) / 20.0, 1.0),
            }

        # Persist to database
        success_rate = len(successful_waits) / len(attempts) if attempts else 0

        # Skip database operations in programmatic mode
        if self.programmatic_mode:
            logger.debug(
                f"Skipping estimate persistence in programmatic mode for {engine_type}"
            )
        else:
            # Try to get research context from search tracker

            context = get_search_context()
            username = None
            password = None
            if context is not None:
                username = context.get("username")
                password = context.get("user_password")

            if username and password:
                try:
                    # Use thread-safe metrics writer to save to user's encrypted database
                    from ...database.thread_metrics import metrics_writer

                    # Set password for this thread if not already set
                    metrics_writer.set_user_password(username, password)

                    db_imports = _get_db_imports()
                    RateLimitEstimate = db_imports.get("RateLimitEstimate")

                    session_inner: "Session"
                    with metrics_writer.get_session(username) as session_inner:
                        # Check if estimate exists
                        estimate: Any = (
                            session_inner.query(RateLimitEstimate)
                            .filter_by(engine_type=engine_type)
                            .first()
                        )

                        if estimate:
                            # Update existing estimate
                            estimate.base_wait_seconds = new_base
                            estimate.min_wait_seconds = min_wait
                            estimate.max_wait_seconds = max_wait
                            estimate.last_updated = time.time()
                            estimate.total_attempts = len(attempts)
                            estimate.success_rate = success_rate
                        else:
                            # Create new estimate
                            estimate = RateLimitEstimate(
                                engine_type=engine_type,
                                base_wait_seconds=new_base,
                                min_wait_seconds=min_wait,
                                max_wait_seconds=max_wait,
                                last_updated=time.time(),
                                total_attempts=len(attempts),
                                success_rate=success_rate,
                            )
                            session_inner.add(estimate)

                except Exception:
                    logger.warning("Failed to persist rate limit estimate")
            else:
                logger.debug(
                    "Skipping rate limit estimate save - no user context"
                )

        logger.info(
            f"Updated rate limit for {engine_type}: {new_base:.2f}s "
            f"(success rate: {success_rate:.1%})"
        )

    def _get_in_memory_stats(
        self, engine_type: Optional[str] = None
    ) -> List[Tuple[str, float, float, float, float, int, float]]:
        """Return stats from in-memory caches (no DB access)."""
        with self._cache_lock:
            stats = []
            engines_to_check = (
                [engine_type]
                if engine_type
                else list(self.current_estimates.keys())
            )
            for engine in engines_to_check:
                if engine in self.current_estimates:
                    est = self.current_estimates[engine]
                    attempts = self.recent_attempts.get(engine)
                    attempt_count = len(attempts) if attempts else 0
                    stats.append(
                        (
                            engine,
                            est["base"],
                            est["min"],
                            est["max"],
                            time.time(),
                            attempt_count,
                            est.get("confidence", 0.0),
                        )
                    )
        return stats

    def get_stats(
        self, engine_type: Optional[str] = None
    ) -> List[Tuple[str, float, float, float, float, int, float]]:
        """
        Get statistics for monitoring.

        Args:
            engine_type: Specific engine to get stats for, or None for all

        Returns:
            List of tuples with engine statistics
        """
        # Skip database operations in CI mode
        if is_ci_environment():
            logger.debug("Skipping database stats in CI mode")
            return self._get_in_memory_stats(engine_type)

        # Skip database operations in programmatic mode
        if self.programmatic_mode:
            return self._get_in_memory_stats(engine_type)

        try:
            context = get_search_context()
            if not context:
                return self._get_in_memory_stats(engine_type)

            username = context.get("username")
            password = context.get("user_password")
            if not username or not password:
                return self._get_in_memory_stats(engine_type)

            db_imports = _get_db_imports()
            RateLimitEstimate = db_imports.get("RateLimitEstimate")

            from ...database.thread_metrics import metrics_writer

            metrics_writer.set_user_password(username, password)

            session_est: "Session"
            with metrics_writer.get_session(username) as session_est:
                if engine_type:
                    estimates: list[Any] = (
                        session_est.query(RateLimitEstimate)
                        .filter_by(engine_type=engine_type)
                        .all()
                    )
                else:
                    estimates = (
                        session_est.query(RateLimitEstimate)
                        .order_by(RateLimitEstimate.engine_type)
                        .all()
                    )

                return [
                    (
                        est.engine_type,
                        est.base_wait_seconds,
                        est.min_wait_seconds,
                        est.max_wait_seconds,
                        est.last_updated,
                        est.total_attempts,
                        est.success_rate,
                    )
                    for est in estimates
                ]
        except Exception:
            logger.warning("Failed to get rate limit stats from DB")
            # Return in-memory stats as fallback
            return self._get_in_memory_stats(engine_type)

    def reset_engine(self, engine_type: str) -> None:
        """
        Reset learned values for a specific engine.

        Args:
            engine_type: Engine to reset
        """
        # Always clear from memory first
        with self._cache_lock:
            if engine_type in self.recent_attempts:
                del self.recent_attempts[engine_type]
            if engine_type in self.current_estimates:
                del self.current_estimates[engine_type]

        # Skip database operations in programmatic mode
        if self.programmatic_mode:
            logger.debug(
                f"Reset rate limit data for {engine_type} (memory only in programmatic mode)"
            )
            return

        # Skip database operations in CI mode
        if is_ci_environment():
            logger.debug(
                f"Reset rate limit data for {engine_type} (memory only in CI mode)"
            )
            return

        try:
            context = get_search_context()
            if not context:
                logger.debug(
                    f"Reset rate limit data for {engine_type} (memory only - no user context)"
                )
                return

            username = context.get("username")
            password = context.get("user_password")
            if not username or not password:
                logger.debug(
                    f"Reset rate limit data for {engine_type} (memory only - no credentials)"
                )
                return

            db_imports = _get_db_imports()
            RateLimitAttempt = db_imports.get("RateLimitAttempt")
            RateLimitEstimate = db_imports.get("RateLimitEstimate")

            from ...database.thread_metrics import metrics_writer

            metrics_writer.set_user_password(username, password)

            session_reset: "Session"
            with metrics_writer.get_session(username) as session_reset:
                # Delete historical attempts
                session_reset.query(RateLimitAttempt).filter_by(
                    engine_type=engine_type
                ).delete()

                # Delete estimates
                session_reset.query(RateLimitEstimate).filter_by(
                    engine_type=engine_type
                ).delete()

                session_reset.commit()

            logger.info(f"Reset rate limit data for {engine_type}")

        except Exception:
            logger.warning(
                f"Failed to reset rate limit data in database for {engine_type}"
                "In-memory data was cleared successfully."
            )
            # Don't re-raise in test contexts - the memory cleanup is sufficient

    def get_search_quality_stats(
        self, engine_type: Optional[str] = None
    ) -> List[Dict]:
        """
        Get basic search quality statistics for monitoring.

        Note: no HTTP route consumes this method any more —
        ``/benchmark/api/search-quality`` now reports success_rate from the
        persisted ``RateLimitEstimate`` table instead of this in-memory,
        per-process view. It is retained as a tracker diagnostic capability
        (with its own test coverage) for CLI / programmatic use; the
        ``recent_avg_results`` / ``sample_size`` shape it returns is no longer
        exposed over the API.

        Args:
            engine_type: Specific engine to get stats for, or None for all

        Returns:
            List of dictionaries with search quality metrics
        """
        stats = []

        with self._cache_lock:
            engines_to_check = (
                [engine_type]
                if engine_type
                else list(self.recent_attempts.keys())
            )
            # Snapshot the data under lock, then process outside
            engine_attempts = {}
            for engine in engines_to_check:
                if engine in self.recent_attempts:
                    engine_attempts[engine] = list(self.recent_attempts[engine])

        for engine, recent in engine_attempts.items():
            search_counts = [
                attempt.get("search_result_count", 0)
                for attempt in recent
                if attempt.get("search_result_count") is not None
            ]

            if not search_counts:
                continue

            recent_avg = sum(search_counts) / len(search_counts)

            stats.append(
                {
                    "engine_type": engine,
                    "recent_avg_results": recent_avg,
                    "min_recent_results": min(search_counts),
                    "max_recent_results": max(search_counts),
                    "sample_size": len(search_counts),
                    "total_attempts": len(recent),
                    "status": self._get_quality_status(recent_avg),
                }
            )

        return stats

    def _get_quality_status(self, recent_avg: float) -> str:
        """Get quality status string based on average results."""
        if recent_avg < 1:
            return "CRITICAL"
        if recent_avg < 3:
            return "WARNING"
        if recent_avg < 5:
            return "CAUTION"
        if recent_avg >= 10:
            return "EXCELLENT"
        return "GOOD"

    def cleanup_old_data(self, days: int = 30) -> None:
        """
        Remove old retry attempt data to prevent database bloat.

        Args:
            days: Remove data older than this many days
        """
        cutoff_time = time.time() - (days * 24 * 3600)

        # Skip database operations in programmatic mode
        if self.programmatic_mode:
            logger.debug("Skipping database cleanup in programmatic mode")
            return

        # Skip database operations in CI mode
        if is_ci_environment():
            logger.debug("Skipping database cleanup in CI mode")
            return

        try:
            context = get_search_context()
            if not context:
                logger.debug("Skipping database cleanup - no user context")
                return

            username = context.get("username")
            password = context.get("user_password")
            if not username or not password:
                logger.debug("Skipping database cleanup - no credentials")
                return

            db_imports = _get_db_imports()
            RateLimitAttempt = db_imports.get("RateLimitAttempt")

            from ...database.thread_metrics import metrics_writer

            metrics_writer.set_user_password(username, password)

            session_clean: "Session"
            with metrics_writer.get_session(username) as session_clean:
                # Count and delete old attempts
                old_attempts: Any = session_clean.query(
                    RateLimitAttempt
                ).filter(RateLimitAttempt.timestamp < cutoff_time)
                deleted_count = old_attempts.count()
                old_attempts.delete()

                session_clean.commit()

            if deleted_count > 0:
                logger.info(f"Cleaned up {deleted_count} old retry attempts")

        except Exception:
            logger.warning("Failed to cleanup old rate limit data")


def get_tracker() -> AdaptiveRateLimitTracker:
    """Create a fresh rate limit tracker instance.

    Returns a new instance each call so that per-user DB credentials
    are never cached across requests in a multi-user Flask app.
    """
    return AdaptiveRateLimitTracker()
