"""Pin ``ProxyUpdateSpend`` behavior.

Symbols pinned here:
  - ``ProxyUpdateSpend.update_end_user_spend``
  - ``ProxyUpdateSpend.update_spend_logs``
  - ``ProxyUpdateSpend.disable_spend_updates``
"""

from __future__ import annotations

import asyncio
from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock

import pytest

from litellm.proxy.utils import ProxyUpdateSpend


class _AsyncCM:
    def __init__(self, target: Any) -> None:
        self.target = target

    async def __aenter__(self) -> Any:
        return self.target

    async def __aexit__(self, *exc: Any) -> None:
        return None


@pytest.mark.asyncio
async def test_update_end_user_spend_upserts_each_end_user(
    mock_prisma_client: Any,
) -> None:
    batcher = MagicMock()
    batcher.litellm_endusertable.upsert = MagicMock()
    transaction = MagicMock()
    transaction.batch_ = lambda: _AsyncCM(batcher)
    mock_prisma_client.db.tx = lambda timeout: _AsyncCM(transaction)

    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()

    end_user_costs: Dict[str, float] = {"u_b": 1.0, "u_a": 0.5}
    await ProxyUpdateSpend.update_end_user_spend(
        n_retry_times=0,
        prisma_client=mock_prisma_client,
        proxy_logging_obj=proxy_logging,
        end_user_list_transactions=end_user_costs,
    )
    calls = batcher.litellm_endusertable.upsert.call_args_list
    ordered_ids = [c.kwargs["where"]["user_id"] for c in calls]
    creates = [c.kwargs["data"]["create"] for c in calls]
    pinned = {
        "upsert_count": len(calls),
        "ordered_ids": ordered_ids,
        "first_create_keys": sorted(creates[0].keys()),
        "first_create_user_id": creates[0]["user_id"],
        "first_create_spend": creates[0]["spend"],
    }
    assert pinned == {
        "upsert_count": 2,
        "ordered_ids": ["u_a", "u_b"],
        "first_create_keys": sorted(["user_id", "spend", "blocked"]),
        "first_create_user_id": "u_a",
        "first_create_spend": 0.5,
    }


@pytest.mark.asyncio
async def test_update_end_user_spend_retries_on_connection_error(
    mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
    """``DB_CONNECTION_ERROR_TYPES`` failures should be retried with backoff;
    once retries are exhausted, ``_raise_failed_update_spend_exception`` is
    invoked and the original exception bubbles up.
    """
    import httpx
    import litellm.proxy.utils as utils_mod

    sleeps: list[float] = []

    async def _fake_sleep(seconds: float) -> None:
        sleeps.append(seconds)

    monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)

    err = httpx.ReadError("conn reset")
    mock_prisma_client.db.tx = MagicMock(side_effect=err)
    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()
    with pytest.raises(httpx.ReadError):
        await ProxyUpdateSpend.update_end_user_spend(
            n_retry_times=1,
            prisma_client=mock_prisma_client,
            proxy_logging_obj=proxy_logging,
            end_user_list_transactions={"u": 1.0},
        )
    assert sleeps == [1.0]


@pytest.mark.asyncio
async def test_update_end_user_spend_non_connection_error_raises_immediately(
    mock_prisma_client: Any,
) -> None:
    mock_prisma_client.db.tx = MagicMock(side_effect=RuntimeError("unknown"))
    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()
    with pytest.raises(RuntimeError, match="unknown"):
        await ProxyUpdateSpend.update_end_user_spend(
            n_retry_times=3,
            prisma_client=mock_prisma_client,
            proxy_logging_obj=proxy_logging,
            end_user_list_transactions={"u": 1.0},
        )


@pytest.mark.asyncio
async def test_update_spend_logs_writes_batches_via_create_many(
    mock_prisma_client: Any, make_spend_log_row: Any
) -> None:
    logs = [make_spend_log_row(request_id=f"r{i}", spend=float(i)) for i in range(3)]
    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock()
    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()
    await ProxyUpdateSpend.update_spend_logs(
        n_retry_times=0,
        prisma_client=mock_prisma_client,
        db_writer_client=None,
        proxy_logging_obj=proxy_logging,
        logs_to_process=logs,
    )
    kwargs = mock_prisma_client.db.litellm_spendlogs.create_many.await_args.kwargs
    pinned = {
        "calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count,
        "data_len": len(kwargs["data"]),
        "skip_duplicates": kwargs["skip_duplicates"],
        "first_request_id": kwargs["data"][0]["request_id"],
    }
    assert pinned == {
        "calls": 1,
        "data_len": 3,
        "skip_duplicates": True,
        "first_request_id": "r0",
    }


@pytest.mark.asyncio
async def test_update_spend_logs_uses_spend_logs_url_when_set(
    mock_prisma_client: Any,
    make_spend_log_row: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("SPEND_LOGS_URL", "http://writer.invalid")
    writer = MagicMock()
    writer.post = AsyncMock(return_value=MagicMock(status_code=200))
    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()
    logs = [make_spend_log_row(request_id="r1")]
    await ProxyUpdateSpend.update_spend_logs(
        n_retry_times=0,
        prisma_client=mock_prisma_client,
        db_writer_client=writer,
        proxy_logging_obj=proxy_logging,
        logs_to_process=logs,
    )
    pinned = {
        "post_calls": writer.post.await_count,
        "url": writer.post.await_args.kwargs["url"],
        "headers": writer.post.await_args.kwargs["headers"],
        "create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count,
    }
    assert pinned == {
        "post_calls": 1,
        "url": "http://writer.invalid/spend/update",
        "headers": {"Content-Type": "application/json"},
        "create_many_calls": 0,
    }


@pytest.mark.asyncio
async def test_update_spend_logs_pops_logs_when_logs_to_process_is_none(
    mock_prisma_client: Any, make_spend_log_row: Any
) -> None:
    mock_prisma_client.spend_log_transactions = [
        make_spend_log_row(request_id="a"),
        make_spend_log_row(request_id="b"),
    ]
    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock()
    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()
    await ProxyUpdateSpend.update_spend_logs(
        n_retry_times=0,
        prisma_client=mock_prisma_client,
        db_writer_client=None,
        proxy_logging_obj=proxy_logging,
    )
    assert mock_prisma_client.spend_log_transactions == []
    assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1


@pytest.mark.asyncio
async def test_update_spend_logs_failure_raises_after_retries(
    mock_prisma_client: Any,
    make_spend_log_row: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """When all retries exhaust the underlying DB error, the helper raises
    via ``_raise_failed_update_spend_exception``.
    """
    import httpx
    import litellm.proxy.utils as utils_mod

    async def _fake_sleep(_: float) -> None:
        return None

    monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)

    mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
        side_effect=httpx.ReadError("network blip")
    )
    proxy_logging = MagicMock()
    proxy_logging.failure_handler = AsyncMock()
    with pytest.raises(httpx.ReadError):
        await ProxyUpdateSpend.update_spend_logs(
            n_retry_times=1,
            prisma_client=mock_prisma_client,
            db_writer_client=None,
            proxy_logging_obj=proxy_logging,
            logs_to_process=[make_spend_log_row(request_id="r1")],
        )


def test_disable_spend_updates_reflects_general_settings(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The static method delegates to ``general_settings['disable_spend_updates']``;
    flipping that value toggles the helper's return.
    """
    import litellm.proxy.proxy_server as proxy_server_mod

    monkeypatch.setattr(
        proxy_server_mod, "general_settings", {"disable_spend_updates": True}
    )
    pinned = {
        "with_flag_true": ProxyUpdateSpend.disable_spend_updates(),
        "type_is_bool": isinstance(ProxyUpdateSpend.disable_spend_updates(), bool),
        "method_is_static": isinstance(
            ProxyUpdateSpend.__dict__["disable_spend_updates"], staticmethod
        ),
    }
    assert pinned == {
        "with_flag_true": True,
        "type_is_bool": True,
        "method_is_static": True,
    }


def test_disable_spend_updates_default_false_without_flag(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    import litellm.proxy.proxy_server as proxy_server_mod

    monkeypatch.setattr(proxy_server_mod, "general_settings", {})
    assert ProxyUpdateSpend.disable_spend_updates() is False


def test_disable_spend_updates_error_when_general_settings_unavailable(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    import litellm.proxy.proxy_server as proxy_server_mod

    monkeypatch.delattr(proxy_server_mod, "general_settings", raising=False)
    with pytest.raises(ImportError):
        ProxyUpdateSpend.disable_spend_updates()
