from __future__ import annotations

import pytest
import respx

from rungate import AsyncRungate, Rungate

MOCK_API_KEY = "rg_sk_test_key_for_unit_tests"
MOCK_BASE_URL = "http://test-rungate:8080"


def make_openai_response(
    *,
    model: str = "gpt-4o",
    content: str = "Hello!",
    tokens: dict | None = None,
) -> dict:
    """Create a standard OpenAI chat completion response body."""
    usage = tokens or {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
    return {
        "id": "chatcmpl-test-123",
        "object": "chat.completion",
        "model": model,
        "created": 1700000000,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": content},
                "finish_reason": "stop",
            }
        ],
        "usage": usage,
    }


def make_rungate_headers(
    *,
    run_id: str = "run-test-123",
    step_id: str = "step-test-456",
    provider: str = "openai",
    model: str = "gpt-4o",
    tokens_used: str = "15",
    cost_usd: str = "0.000075",
    budget_remaining_usd: str | None = None,
    budget_remaining_tokens: str | None = None,
    failover: str | None = None,
    warning: str | None = None,
    governance: str | None = None,
) -> dict[str, str]:
    """Create standard Rungate response headers."""
    headers: dict[str, str] = {
        "x-rungate-run-id": run_id,
        "x-rungate-step-id": step_id,
        "x-rungate-provider": provider,
        "x-rungate-model": model,
        "x-rungate-tokens-used": tokens_used,
        "x-rungate-cost-usd": cost_usd,
    }
    if budget_remaining_usd is not None:
        headers["x-rungate-budget-remaining-usd"] = budget_remaining_usd
    if budget_remaining_tokens is not None:
        headers["x-rungate-budget-remaining-tokens"] = budget_remaining_tokens
    if failover is not None:
        headers["x-rungate-failover"] = failover
    if warning is not None:
        headers["x-rungate-warning"] = warning
    if governance is not None:
        headers["x-rungate-governance"] = governance
    return headers


@pytest.fixture
def mock_router():
    """respx mock router for intercepting httpx requests."""
    with respx.mock(assert_all_called=False) as router:
        yield router


@pytest.fixture
def client():
    """Synchronous Rungate client for testing."""
    c = Rungate(api_key=MOCK_API_KEY, base_url=MOCK_BASE_URL)
    yield c
    c.close()


@pytest.fixture
async def async_client():
    """Async Rungate client for testing."""
    c = AsyncRungate(api_key=MOCK_API_KEY, base_url=MOCK_BASE_URL)
    yield c
    await c.close()
