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

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

from loguru import logger

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

# 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,
                )
                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)

        # In-memory cache for fast access
        self.recent_attempts: Dict[str, deque] = {}
        self.current_estimates: Dict[str, Dict[str, float]] = {}

        # 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

        # Skip database operations in fallback mode
        if use_fallback_llm():
            logger.debug(
                "Skipping rate limit estimate loading in fallback mode"
            )
            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)

                with metrics_writer.get_session(username) as session:
                    estimates = 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)

                        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 as e:
                logger.warning(f"Could not load rate limit estimates: {e}")
                # 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()

        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 = {
                "LocalSearchEngine": 0.0,  # No network calls
                "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"]

        # 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(estimate["min"], min(wait_time, estimate["max"]))
        return 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
        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."""
        if (
            engine_type not in self.recent_attempts
            or len(self.recent_attempts[engine_type]) < 3
        ):
            logger.info(
                f"Not updating estimate for {engine_type} - only {len(self.recent_attempts.get(engine_type, []))} attempts (need 3)"
            )
            return

        attempts = list(self.recent_attempts[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 engine_type in self.current_estimates:
            old_base = self.current_estimates[engine_type]["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
        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 (skip in fallback mode)
        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}"
            )
        elif not use_fallback_llm():
            # 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")

                    with metrics_writer.get_session(username) as session:
                        # Check if estimate exists
                        estimate = (
                            session.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.add(estimate)

                except Exception as e:
                    logger.exception(
                        f"Failed to persist rate limit estimate: {e}"
                    )
            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_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 test/fallback mode
        if use_fallback_llm() or is_ci_environment():
            logger.debug("Skipping database stats in test/CI mode")
            # Return stats from in-memory estimates
            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]
                    stats.append(
                        (
                            engine,
                            est["base"],
                            est["min"],
                            est["max"],
                            time.time(),
                            len(self.recent_attempts.get(engine, [])),
                            est.get("confidence", 0.0),
                        )
                    )
            return stats

        # Skip database operations in programmatic mode
        if self.programmatic_mode:
            return stats

        try:
            db_imports = _get_db_imports()
            get_user_db_session = db_imports.get("get_user_db_session")
            RateLimitEstimate = db_imports.get("RateLimitEstimate")

            with get_user_db_session() as session:
                if engine_type:
                    estimates = (
                        session.query(RateLimitEstimate)
                        .filter_by(engine_type=engine_type)
                        .all()
                    )
                else:
                    estimates = (
                        session.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 as e:
            logger.warning(f"Failed to get rate limit stats from DB: {e}")
            # Return in-memory stats as fallback
            return self.get_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
        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 test/fallback mode
        if use_fallback_llm() or is_ci_environment():
            logger.debug(
                f"Reset rate limit data for {engine_type} (memory only in test/CI mode)"
            )
            return

        try:
            db_imports = _get_db_imports()
            get_user_db_session = db_imports.get("get_user_db_session")
            RateLimitAttempt = db_imports.get("RateLimitAttempt")
            RateLimitEstimate = db_imports.get("RateLimitEstimate")

            with get_user_db_session() as session:
                # Delete historical attempts
                session.query(RateLimitAttempt).filter_by(
                    engine_type=engine_type
                ).delete()

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

                session.commit()

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

        except Exception as e:
            logger.warning(
                f"Failed to reset rate limit data in database for {engine_type}: {e}. "
                "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.

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

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

        engines_to_check = (
            [engine_type] if engine_type else list(self.recent_attempts.keys())
        )

        for engine in engines_to_check:
            if engine not in self.recent_attempts:
                continue

            recent = list(self.recent_attempts[engine])
            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"
        elif recent_avg < 3:
            return "WARNING"
        elif recent_avg < 5:
            return "CAUTION"
        elif recent_avg >= 10:
            return "EXCELLENT"
        else:
            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 test/fallback mode
        if use_fallback_llm() or is_ci_environment():
            logger.debug("Skipping database cleanup in test/CI mode")
            return

        try:
            db_imports = _get_db_imports()
            get_user_db_session = db_imports.get("get_user_db_session")
            RateLimitAttempt = db_imports.get("RateLimitAttempt")

            with get_user_db_session() as session:
                # Count and delete old attempts
                old_attempts = session.query(RateLimitAttempt).filter(
                    RateLimitAttempt.timestamp < cutoff_time
                )
                deleted_count = old_attempts.count()
                old_attempts.delete()

                session.commit()

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

        except Exception as e:
            logger.warning(f"Failed to cleanup old rate limit data: {e}")


# Create a singleton instance
_tracker_instance: Optional[AdaptiveRateLimitTracker] = None


def get_tracker() -> AdaptiveRateLimitTracker:
    """Get the global rate limit tracker instance."""
    global _tracker_instance
    if _tracker_instance is None:
        _tracker_instance = AdaptiveRateLimitTracker()
    return _tracker_instance
