from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from functools import wraps
from typing import Any, Callable, TypeVar

from src.core.exceptions import ConnectorException
from src.core.logging import get_logger

logger = get_logger("connector_retry")

T = TypeVar("T")

DEFAULT_RETRYABLE_EXCEPTIONS = (ConnectionError, TimeoutError, OSError)


@dataclass
class RetryPolicy:
    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 30.0
    exponential_base: float = 2.0
    jitter: bool = True
    retryable_exceptions: tuple[type[Exception], ...] = field(default_factory=lambda: DEFAULT_RETRYABLE_EXCEPTIONS)

    def delay_for_attempt(self, attempt: int) -> float:
        import random

        delay = min(self.base_delay * (self.exponential_base ** (attempt - 1)), self.max_delay)
        if self.jitter:
            delay = delay * (0.5 + random.random() * 0.5)
        return delay

    def should_retry(self, exception: Exception) -> bool:
        if isinstance(exception, ConnectorException):
            return False
        return isinstance(exception, self.retryable_exceptions)


async def with_retry(
    func: Callable[..., Any],
    *args: Any,
    policy: RetryPolicy | None = None,
    connector_type: str = "unknown",
    operation: str = "unknown",
    **kwargs: Any,
) -> Any:
    p = policy or RetryPolicy()
    last_exception: Exception | None = None

    for attempt in range(1, p.max_attempts + 1):
        try:
            result = await func(*args, **kwargs)
            if attempt > 1:
                logger.info("retry_succeeded", extra={"connector_type": connector_type, "operation": operation, "attempt": attempt})
            return result
        except Exception as exc:
            last_exception = exc
            if not p.should_retry(exc) or attempt >= p.max_attempts:
                break
            delay = p.delay_for_attempt(attempt)
            logger.warning(
                "retry_attempt",
                extra={
                    "connector_type": connector_type,
                    "operation": operation,
                    "attempt": attempt,
                    "delay_s": round(delay, 2),
                    "error": str(exc),
                },
            )
            await asyncio.sleep(delay)

    raise last_exception  # type: ignore[misc]


def retry(policy: RetryPolicy | None = None):
    p = policy or RetryPolicy()

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            connector_type = getattr(args[0], "connector_type", "unknown") if args else "unknown"
            return await with_retry(func, *args, policy=p, connector_type=connector_type, operation=func.__name__, **kwargs)
        return wrapper
    return decorator