"""Tests for the planned RetryPolicy-based LLM client API."""

from __future__ import annotations

import types

import pytest

from agentbyte.component import ComponentLoader
from agentbyte.llm import (
    AzureOpenAIChatCompletionClient,
    AzureOpenAIEmbeddingClient,
    OpenAIChatCompletionClient,
    OpenAIEmbeddingClient,
    RetryPolicy,
)


class _FakeOpenAIClient:
    def __init__(self) -> None:
        self.chat = types.SimpleNamespace(
            completions=types.SimpleNamespace(create=self._create_chat)
        )
        self.embeddings = types.SimpleNamespace(create=self._create_embedding)

    async def _create_chat(self, **kwargs):
        return kwargs

    async def _create_embedding(self, **kwargs):
        return kwargs


def test_retry_policy_defaults() -> None:
    policy = RetryPolicy()

    assert policy.max_retries == 3
    assert policy.initial_retry_delay == 1.0
    assert policy.max_retry_delay == 60.0


def test_retry_policy_rejects_negative_delay() -> None:
    with pytest.raises(ValueError):
        RetryPolicy(initial_retry_delay=-1.0)


@pytest.mark.parametrize(
    ("client_cls", "model"),
    [
        (OpenAIChatCompletionClient, "gpt-4.1-mini"),
        (AzureOpenAIChatCompletionClient, "gpt-4.1-mini"),
        (OpenAIEmbeddingClient, "text-embedding-3-small"),
        (AzureOpenAIEmbeddingClient, "embed-prod"),
    ],
)
def test_clients_accept_retry_policy_constructor_argument(client_cls, model) -> None:
    policy = RetryPolicy(max_retries=5, initial_retry_delay=0.25, max_retry_delay=10.0)

    client = client_cls(
        model=model,
        client=_FakeOpenAIClient(),
        retry_policy=policy,
    )

    assert client.retry_policy == policy


@pytest.mark.parametrize(
    ("client_cls", "model"),
    [
        (OpenAIChatCompletionClient, "gpt-4.1-mini"),
        (AzureOpenAIChatCompletionClient, "gpt-4.1-mini"),
        (OpenAIEmbeddingClient, "text-embedding-3-small"),
        (AzureOpenAIEmbeddingClient, "embed-prod"),
    ],
)
def test_component_dump_includes_nested_retry_policy(client_cls, model) -> None:
    policy = RetryPolicy(max_retries=4, initial_retry_delay=0.5, max_retry_delay=8.0)

    client = client_cls(
        model=model,
        client=_FakeOpenAIClient(),
        retry_policy=policy,
    )

    dumped = client.dump_component()

    assert dumped.config["retry_policy"] == {
        "max_retries": 4,
        "initial_retry_delay": 0.5,
        "max_retry_delay": 8.0,
    }
    assert "max_retries" not in dumped.config
    assert "initial_retry_delay" not in dumped.config
    assert "max_retry_delay" not in dumped.config


def test_openai_chat_component_load_restores_retry_policy(monkeypatch) -> None:
    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")

    component = {
        "provider": "agentbyte.llm.OpenAIChatCompletionClient",
        "component_type": "model_client",
        "config": {
            "model": "gpt-4.1-mini",
            "config": {},
            "key_env_var": "OPENAI_API_KEY",
            "retry_policy": {
                "max_retries": 7,
                "initial_retry_delay": 0.2,
                "max_retry_delay": 9.0,
            },
        },
    }

    restored = ComponentLoader.load_component(component)

    assert isinstance(restored, OpenAIChatCompletionClient)
    assert restored.retry_policy.max_retries == 7
    assert restored.retry_policy.initial_retry_delay == 0.2
    assert restored.retry_policy.max_retry_delay == 9.0
