"""Adaptive concurrency control for LLM API calls.

This module provides a feedback-driven concurrency controller that discovers
effective limits by observing outcomes (success, rate-limit, error). It uses
the AIMD (Additive Increase, Multiplicative Decrease) algorithm from TCP
congestion control.

Key concepts:
- **ConcurrencyKey**: Determines the isolation boundary for rate limits.
  Requests with different keys have independent limits.
- **Outcome**: The result of a provider call (SUCCESS, RATE_LIMITED, ERROR).
- **AdaptiveConcurrencyLimiter**: The core AIMD algorithm.
- **with_limit**: Context manager for ergonomic usage at call sites.

Independence is emergent from keying: by default, we key by provider name
(e.g., "openai", "anthropic"), which matches how most providers enforce limits.

Example usage:
    from ember._internal.concurrency import (
        AdaptiveConcurrencyLimiter,
        ConcurrencyKey,
        Outcome,
        with_limit,
    )

    # Simple usage with context manager (recommended):
    with with_limit("openai/gpt-4o"):
        response = provider.complete(...)

    # Manual usage for more control:
    limiter = AdaptiveConcurrencyLimiter()
    key = ConcurrencyKey(provider="openai")

    acquired = limiter.acquire(key, timeout_s=30.0)
    if not acquired:
        raise TimeoutError("Concurrency limit timeout")

    try:
        response = call_provider(...)
        limiter.release(key, Outcome.SUCCESS)
    except RateLimitError:
        limiter.release(key, Outcome.RATE_LIMITED)
        raise

References:
    - Netflix concurrency-limits: https://github.com/Netflix/concurrency-limits
    - TCP Congestion Control (RFC 5681): https://datatracker.ietf.org/doc/html/rfc5681
"""

from __future__ import annotations

from contextlib import contextmanager
from typing import Generator

from ember._internal.concurrency.limiter import (
    AdaptiveConcurrencyLimiter,
    AdjustmentCallback,
    ConcurrencyKey,
    Outcome,
    get_limiter,
    outcome_from_error,
)


@contextmanager
def with_limit(
    model: str,
    *,
    timeout_s: float = 30.0,
) -> Generator[None, None, None]:
    """Context manager for adaptive concurrency control keyed by model's provider.

    Uses AIMD (TCP-style) algorithm to discover effective concurrency limits.
    On 429 rate limit: limit halves. On success: limit slowly increases.

    Example:
        with with_limit("openai/gpt-4o"):
            response = provider.complete(...)

    Args:
        model: Model identifier (e.g., "openai/gpt-4o", "gpt-5-codex").
        timeout_s: Maximum seconds to wait for a slot. Default 30.

    Raises:
        ProviderAPIError: If timeout waiting for capacity.
    """
    from ember._internal.exceptions import ProviderAPIError

    limiter = get_limiter()
    key = ConcurrencyKey.from_model_id(model)

    if not limiter.acquire(key, timeout_s=timeout_s):
        raise ProviderAPIError(
            "Concurrency limit timeout waiting for provider capacity",
            context={"model": model, "provider": key.provider, "timeout_s": timeout_s},
        )

    try:
        yield
    except Exception as exc:
        limiter.release(key, outcome_from_error(exc))
        raise
    else:
        limiter.release(key, Outcome.SUCCESS)


__all__ = [
    "AdaptiveConcurrencyLimiter",
    "AdjustmentCallback",
    "ConcurrencyKey",
    "Outcome",
    "with_limit",
]
