import logging
import random
import datetime
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    RetryError,
    RetryCallState,
    wait_base, # Import for custom wait strategies
)
from typing import Any, Callable, TypeVar, Optional, Union

from .config import RetryConfig
# Import the transport_registry instance instead of the old function
from .transports import Transport, transport_registry

logger = logging.getLogger(__name__)

# The type that the user's callable_action is expected to return upon success.
# This is NOT what attempt_action returns, as attempt_action can return an Exception for retry state.
ExpectedResponseType = TypeVar("ExpectedResponseType")


class WaitWith429AndJitter(wait_base):
    def __init__(
        self,
        base_wait_strategy: wait_base,
        config: RetryConfig,
        max_jitter: float = 1.0, # Default max jitter in seconds
    ):
        self.base_wait_strategy = base_wait_strategy
        self.config = config
        self.max_jitter = max_jitter

    def __call__(self, retry_state: RetryCallState) -> float:
        last_attempt_result = None
        if not retry_state.outcome.failed:
            last_attempt_result = retry_state.outcome.result()

        # Check for 429 response first
        if (
            self.config.handle_429
            and hasattr(last_attempt_result, "status_code")
            and last_attempt_result.status_code == 429
            and hasattr(last_attempt_result, "headers")
        ):
            retry_after = last_attempt_result.headers.get("Retry-After")
            if retry_after:
                try:
                    # RFC 7231 allows seconds (integer) or a HTTP-date
                    wait_time = float(retry_after)
                    logger.info(f"429: Waiting for Retry-After: {wait_time:.2f} seconds.")
                    return wait_time
                except ValueError:
                    try:
                        # Attempt to parse as HTTP-date
                        retry_date = datetime.datetime.strptime(
                            retry_after, "%a, %d %b %Y %H:%M:%S GMT"
                        )
                        # Calculate time difference from now
                        now = datetime.datetime.utcnow()
                        wait_time = (retry_date - now).total_seconds()
                        if wait_time > 0:
                            logger.info(f"429: Waiting until {retry_date.isoformat()} ({wait_time:.2f} seconds).")
                            return wait_time
                        logger.warning(f"429: Retry-After date in past or invalid: {retry_after}")
                    except ValueError:
                        logger.warning(f"429: Could not parse Retry-After header: {retry_after}")

            # Fallback for 429 if no valid Retry-After or parsing failed
            calculated_wait = self.config.default_429_wait_time
            logger.info(f"429: Using default wait time: {calculated_wait:.2f} seconds.")
            return calculated_wait

        # If not a 429, or 429 handling is off, use the base strategy (e.g., exponential backoff)
        calculated_wait = self.base_wait_strategy(retry_state)

        # Add jitter if configured
        if self.config.add_jitter:
            jitter = random.uniform(0, self.max_jitter)
            calculated_wait += jitter
            logger.debug(f"Adding {jitter:.2f}s jitter. New wait: {calculated_wait:.2f}s")

        return calculated_wait


def _is_response_valid(response: Any, config: RetryConfig) -> bool:
    """Checks if a given response object meets the configured criteria."""
    # Ensure it at least looks like an HTTP response object
    if not hasattr(response, "status_code") or not callable(
        getattr(response, "json", None)
    ):
        logger.warning(
            f"Validation Failed: Object {type(response).__name__} lacks "
            f"expected 'status_code' or 'json()' method."
        )
        return False

    # If 429 handling is on, treat 429 as *invalid* for _is_response_valid's purpose
    # so that should_retry specifically picks it up for its custom wait logic.
    if config.handle_429 and response.status_code == 429:
        logger.debug(f"Response status 429 being handled by custom wait strategy.")
        return False # This will trigger retry as `not is_valid` in should_retry

    if config.status_codes_set:
        if response.status_code not in config.status_codes_set:
            logger.warning(
                f"Retry: Invalid status {response.status_code} (expected one of {config.status_codes_set})"
            )
            return False
    # General check for 2xx codes (response.ok for requests, response.is_success for httpx, or direct 2xx check)
    elif not getattr(response, "is_success", None) and not getattr(
        response, "ok", None
    ):
        if not (200 <= response.status_code < 300):
            logger.warning(f"Retry: Non-2xx status {response.status_code}")
            return False

    if config.check_json:
        try:
            data = response.json()
        except Exception:  # Catch JSONDecodeError, simplejson.errors.JSONDecodeError, etc.
            logger.warning("Retry: Response is not valid JSON.")
            return False

        if config.required_json_keys:
            missing = [k for k in config.required_json_keys if k not in data]
            if missing:
                logger.warning(f"Retry: Missing JSON keys: {missing}")
                return False
    return True


def invoke(
    callable_action: Callable[[], ExpectedResponseType],
    config: RetryConfig,
    transport: Optional[Transport] = None,
) -> ExpectedResponseType:
    """
    Executes a callable with robust retry and validation logic.

    Args:
        callable_action: A zero-argument function that performs an action,
            like `lambda: requests.get(url)`.
        config: A `RetryConfig` object with validation and retry rules.
        transport: A `Transport` object that defines the library-specific
            behavior. If None, a default is chosen (requests > httpx).

    Returns:
        The validated result from the callable_action.

    Raises:
        tenacity.RetryError: If all attempts fail.
        Any other exception: If a non-retryable exception occurs in callable_action.
    """
    active_transport = transport or transport_registry.get_default()
    logger.debug(f"Using transport: {active_transport.name}")

    def should_retry(retry_state: RetryCallState) -> bool:
        """Determines if tenacity should retry based on RetryCallState."""
        if retry_state.outcome.failed:
            exc = retry_state.outcome.exception()
            # 1. If it's a transport-related exception (connection, timeout, etc.), retry.
            if isinstance(exc, active_transport.exception_type):
                logger.warning(
                    f"Retry: Action failed with transport exception: {type(exc).__name__} - {exc}"
                )
                return True
            # 2. For any other kind of exception (e.g., a TypeError in the callable itself),
            #    we do not retry. Let tenacity re-raise it immediately via `reraise=True`.
            logger.error(
                f"Non-transport exception occurred: {type(exc).__name__} - {exc}. Not retrying."
            )
            return False
        else:  # Outcome was successful (returned a value)
            result = retry_state.outcome.result()
            # 3. Check if the returned result is valid based on config.
            #    _is_response_valid will now do duck-typing check internally.
            #    If _is_response_valid fails, it means the response is not as expected, so retry.
            is_valid = _is_response_valid(result, config)
            if not is_valid:
                logger.warning(
                    f"Retry: Response validation failed for {type(result).__name__}."
                )
                return True # Not valid -> retry

            # If all checks pass, we don't need to retry.
            return False

    # Define the base wait strategy (exponential backoff)
    base_wait = wait_exponential(multiplier=config.backoff_factor, min=1, max=10)

    # Wrap the base wait strategy with our custom 429/jitter logic
    custom_wait_strategy = WaitWith429AndJitter(
        base_wait_strategy=base_wait,
        config=config,
        max_jitter=config.default_429_wait_time # Using default_429_wait_time as max jitter range
    )

    @retry(
        stop=stop_after_attempt(config.max_retries),
        wait=custom_wait_strategy, # Use our custom wait strategy here
        retry=should_retry,
        reraise=True,
    )
    def attempt_action() -> ExpectedResponseType:
        """
        Executes the user-provided callable_action.
        tenacity will automatically catch exceptions and store them in outcome.
        """
        return callable_action()

    return attempt_action()


# Export the public API
from .config import RetryConfig
from .transports import (
    Transport,
    transport_registry,
    REQUESTS_TRANSPORT,
    HTTPX_TRANSPORT,
    NIQUESTS_TRANSPORT,
)

__all__ = [
    "invoke",
    "RetryConfig",
    "RetryError",
    "Transport",
    "transport_registry",
    "REQUESTS_TRANSPORT",
    "HTTPX_TRANSPORT",
    "NIQUESTS_TRANSPORT",
]
