"""
Rate-limited wrapper for LLM calls.
"""

from typing import Optional
from urllib.parse import urlparse

from loguru import logger
from tenacity import (
    retry,
    stop_after_attempt,
    retry_if_exception,
)
from tenacity.wait import wait_base

from ..tracker import get_tracker
from ..exceptions import RateLimitError
from .detection import is_llm_rate_limit_error, extract_retry_after


class AdaptiveLLMWait(wait_base):
    """Adaptive wait strategy for LLM rate limiting."""

    def __init__(self, tracker, engine_type: str):
        self.tracker = tracker
        self.engine_type = engine_type
        self.last_error = None

    def __call__(self, retry_state) -> float:
        # Store the error for potential retry-after extraction
        if retry_state.outcome and retry_state.outcome.failed:
            self.last_error = retry_state.outcome.exception()

        # Get adaptive wait time from tracker
        wait_time = self.tracker.get_wait_time(self.engine_type)

        # If we have a retry-after from the error, use it
        if self.last_error:
            retry_after = extract_retry_after(self.last_error)
            if retry_after > 0:
                wait_time = max(wait_time, retry_after)

        logger.info(
            f"LLM rate limit wait for {self.engine_type}: {wait_time:.2f}s"
        )
        return wait_time


def create_rate_limited_llm_wrapper(base_llm, provider: Optional[str] = None):
    """
    Create a rate-limited wrapper around an LLM instance.

    Args:
        base_llm: The base LLM instance to wrap
        provider: Optional provider name (e.g., 'openai', 'anthropic')

    Returns:
        A wrapped LLM instance with rate limiting capabilities
    """

    class RateLimitedLLMWrapper:
        """Wrapper that adds rate limiting to LLM calls."""

        def __init__(self, llm, provider_name: Optional[str] = None):
            self.base_llm = llm
            self.provider = provider_name
            self.rate_limiter = None

            # Only setup rate limiting if enabled
            if self._should_rate_limit():
                self.rate_limiter = get_tracker()
                logger.info(
                    f"Rate limiting enabled for LLM provider: {self._get_rate_limit_key()}"
                )

        def _should_rate_limit(self) -> bool:
            """Check if rate limiting should be applied to this LLM."""
            # Rate limiting for LLMs is currently disabled by default
            # TODO: Pass settings_snapshot to enable proper configuration
            return False

        def _check_if_local_model(self) -> bool:
            """Check if the LLM is a local model that shouldn't be rate limited."""
            # Don't rate limit local models
            local_providers = [
                "ollama",
                "lmstudio",
                "llamacpp",
                "vllm",
                "local",
                "none",
            ]
            if self.provider and self.provider.lower() in local_providers:
                logger.debug(
                    f"Skipping rate limiting for local provider: {self.provider}"
                )
                return True

            # Check if base URL indicates local model
            if hasattr(self.base_llm, "base_url"):
                base_url = str(self.base_llm.base_url)
                if any(
                    local in base_url
                    for local in ["localhost", "127.0.0.1", "0.0.0.0"]
                ):
                    logger.debug(
                        f"Skipping rate limiting for local URL: {base_url}"
                    )
                    return True

            return False

        def _get_rate_limit_key(self) -> str:
            """Build composite key: provider-url-model"""
            provider = self.provider or "unknown"

            # Extract URL
            url = "unknown"
            if hasattr(self.base_llm, "base_url"):
                url = str(self.base_llm.base_url)
            elif hasattr(self.base_llm, "_client") and hasattr(
                self.base_llm._client, "base_url"
            ):
                url = str(self.base_llm._client.base_url)

            # Clean URL: remove protocol and trailing slashes
            if url != "unknown":
                parsed = urlparse(url)
                url = parsed.netloc or parsed.path
                url = url.rstrip("/")

            # Extract model
            model = "unknown"
            if hasattr(self.base_llm, "model_name"):
                model = str(self.base_llm.model_name)
            elif hasattr(self.base_llm, "model"):
                model = str(self.base_llm.model)

            # Clean model name
            model = model.replace("/", "-").replace(":", "-")

            key = f"{provider}-{url}-{model}"
            return key

        def invoke(self, *args, **kwargs):
            """Invoke the LLM with rate limiting if enabled."""
            if self.rate_limiter:
                rate_limit_key = self._get_rate_limit_key()

                # Define retry logic
                @retry(
                    wait=AdaptiveLLMWait(self.rate_limiter, rate_limit_key),
                    stop=stop_after_attempt(3),
                    retry=retry_if_exception(is_llm_rate_limit_error),
                )
                def _invoke_with_retry():
                    return self._do_invoke(*args, **kwargs)

                try:
                    result = _invoke_with_retry()

                    # Record successful attempt
                    self.rate_limiter.record_outcome(
                        engine_type=rate_limit_key,
                        wait_time=0,  # First attempt had no wait
                        success=True,
                        retry_count=0,
                    )

                    return result

                except Exception as e:
                    # Only record rate limit failures, not general failures
                    if is_llm_rate_limit_error(e):
                        self.rate_limiter.record_outcome(
                            engine_type=rate_limit_key,
                            wait_time=0,
                            success=False,
                            retry_count=0,
                        )
                    raise
            else:
                # No rate limiting, just invoke directly
                return self._do_invoke(*args, **kwargs)

        def _do_invoke(self, *args, **kwargs):
            """Actually invoke the LLM."""
            try:
                return self.base_llm.invoke(*args, **kwargs)
            except Exception as e:
                # Check if it's a rate limit error and wrap it
                if is_llm_rate_limit_error(e):
                    logger.warning(f"LLM rate limit error detected: {e}")
                    raise RateLimitError(f"LLM rate limit: {str(e)}")
                raise

        # Pass through any other attributes to the base LLM
        def __getattr__(self, name):
            return getattr(self.base_llm, name)

        def __str__(self):
            return f"RateLimited({str(self.base_llm)})"

        def __repr__(self):
            return f"RateLimitedLLMWrapper({repr(self.base_llm)})"

    return RateLimitedLLMWrapper(base_llm, provider)
