from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass
from typing import Any, Callable

from .config import RetryConfig
from .exceptions import RetryError
from .utils import compute_delay


@dataclass(slots=True)
class RetryState:
    attempt: int
    max_attempts: int
    delay: float
    exception: BaseException
    function_name: str
    result: Any = None
    elapsed: float = 0.0
    will_retry: bool = True


def build_state(
    *,
    attempt: int,
    config: RetryConfig,
    delay: float,
    exception: BaseException,
    function_name: str,
    result: Any = None,
    elapsed: float = 0.0,
    will_retry: bool = True,
) -> RetryState:
    return RetryState(
        attempt=attempt,
        max_attempts=config.max_attempts,
        delay=delay,
        exception=exception,
        function_name=function_name,
        result=result,
        elapsed=elapsed,
        will_retry=will_retry,
    )


def call_hook(hook: Callable[[RetryState], None] | None, state: RetryState) -> None:
    if hook is not None:
        hook(state)  # HT Check this


def log_retry(config: RetryConfig, state: RetryState) -> None:
    if config.logger is None:
        return

    if state.exception is not None:
        reason = f"exception={type(state.exception).__name__}: {state.exception}"
    else:
        reason = f"retry_if_result matched for result={state.result!r}"

    config.logger.log(
        config.log_level,
        "Retrying %s attempt %s/%s in %.2fs (%s)",
        state.function_name,
        state.attempt,
        state.max_attempts,
        state.delay,
        reason,
    )


def handle_final_failure(
    config: RetryConfig,
    func_name: str,
    last_exception: BaseException,
) -> Any:
    if config.raise_last:
        raise last_exception

    raise RetryError(
        f"Function '{func_name}' failed after {config.max_attempts} attempts",
        last_exception=(
            last_exception if isinstance(last_exception, Exception) else None
        ),
    )


def handle_result_failure(config: RetryConfig, func_name: str, last_result: Any) -> Any:
    raise RetryError(
        f"Function '{func_name}' returned a retryable result after {config.max_attempts} attempts: {last_result!r}"
    )


def should_retry_result(config: RetryConfig, result: Any) -> bool:
    if config.retry_if_result is None:
        return False
    return bool(config.retry_if_result(result))


def maybe_check_timeout(config: RetryConfig, started_at: float, func_name: str) -> None:
    if config.timeout is None:
        return

    elapsed = time.perf_counter() - started_at
    if elapsed > config.timeout:
        raise RetryError(
            f"Function '{func_name}' exceeded retry timeout of {config.timeout:.2f}s"
        )


def run_sync(func, config: RetryConfig, *args, **kwargs):
    last_exception: BaseException | None = None
    last_result: Any = None
    started_at = time.perf_counter()

    for attempt in range(1, config.max_attempts + 1):
        maybe_check_timeout(config, started_at, func.__name__)

        try:
            result = func(*args, **kwargs)
            last_result = result

            if not should_retry_result(config, result):
                return result

            if attempt >= config.max_attempts:
                break

            delay = compute_delay(config, attempt)
            elapsed = time.perf_counter() - started_at
            state = build_state(
                attempt=attempt,
                config=config,
                delay=delay,
                exception=None,
                function_name=func.__name__,
                result=result,
                elapsed=elapsed,
                will_retry=True,
            )
            call_hook(config.before_retry, state)
            call_hook(config.on_retry, state)
            log_retry(config, state)
            time.sleep(delay)
            state.elapsed = time.perf_counter() - started_at
            call_hook(config.after_retry, state)

        except config.exceptions as exc:
            last_exception = exc

            if attempt >= config.max_attempts:
                break

            delay = compute_delay(config, attempt)
            elapsed = time.perf_counter() - started_at
            state = build_state(
                attempt=attempt,
                config=config,
                delay=delay,
                exception=exc,
                function_name=func.__name__,
                result=None,
                elapsed=elapsed,
                will_retry=True,
            )
            call_hook(config.before_retry, state)
            call_hook(config.on_retry, state)
            log_retry(config, state)
            time.sleep(delay)
            state.elapsed = time.perf_counter() - started_at
            call_hook(config.after_retry, state)

    if last_exception is not None:
        return handle_final_failure(config, func.__name__, last_exception)

    return handle_result_failure(config, func.__name__, last_result)


async def run_async(func, config: RetryConfig, *args, **kwargs):
    last_exception: BaseException | None = None
    last_result: Any = None
    started_at = time.perf_counter()

    for attempt in range(1, config.max_attempts + 1):
        maybe_check_timeout(config, started_at, func.__name__)

        try:
            result = await func(*args, **kwargs)
            last_result = result

            if not should_retry_result(config, result):
                return result

            if attempt >= config.max_attempts:
                break

            delay = compute_delay(config, attempt)
            elapsed = time.perf_counter() - started_at
            state = build_state(
                attempt=attempt,
                config=config,
                delay=delay,
                exception=None,
                function_name=func.__name__,
                result=result,
                elapsed=elapsed,
                will_retry=True,
            )
            call_hook(config.before_retry, state)
            call_hook(config.on_retry, state)
            log_retry(config, state)
            await asyncio.sleep(delay)
            state.elapsed = time.perf_counter() - started_at
            call_hook(config.after_retry, state)

        except config.exceptions as exc:
            last_exception = exc

            if attempt >= config.max_attempts:
                break

            delay = compute_delay(config, attempt)
            elapsed = time.perf_counter() - started_at
            state = build_state(
                attempt=attempt,
                config=config,
                delay=delay,
                exception=exc,
                function_name=func.__name__,
                result=None,
                elapsed=elapsed,
                will_retry=True,
            )
            call_hook(config.before_retry, state)
            call_hook(config.on_retry, state)
            log_retry(config, state)
            await asyncio.sleep(delay)
            state.elapsed = time.perf_counter() - started_at
            call_hook(config.after_retry, state)

    if last_exception is not None:
        return handle_final_failure(config, func.__name__, last_exception)

    return handle_result_failure(config, func.__name__, last_result)
