"""Tests for the OpenAI integration."""

import time
from unittest.mock import MagicMock, patch

from agentcost import AgentCost
from agentcost.integrations.openai import track_openai, _TrackedCompletions, _TrackedChat


def _make_ac():
    """Create a debug-mode AgentCost instance."""
    return AgentCost(debug=True)


def _make_mock_response(model="gpt-4o", prompt_tokens=100, completion_tokens=50):
    """Create a mock OpenAI response with usage info."""
    usage = MagicMock()
    usage.prompt_tokens = prompt_tokens
    usage.completion_tokens = completion_tokens

    response = MagicMock()
    response.model = model
    response.usage = usage
    return response


def _make_mock_client(response=None):
    """Create a mock OpenAI client."""
    if response is None:
        response = _make_mock_response()

    client = MagicMock()
    client.chat.completions.create.return_value = response
    return client


class TestTrackOpenAI:
    def test_returns_wrapped_client(self):
        ac = _make_ac()
        client = _make_mock_client()
        tracked = track_openai(client, ac)
        assert tracked is not client
        assert hasattr(tracked, "chat")

    def test_preserves_other_attributes(self):
        ac = _make_ac()
        client = _make_mock_client()
        client.api_key = "sk-test"
        tracked = track_openai(client, ac)
        assert tracked.api_key == "sk-test"

    def test_default_agent_id(self):
        ac = _make_ac()
        client = _make_mock_client()
        tracked = track_openai(client, ac)
        assert tracked.chat.completions._agent_id == "openai-agent"

    def test_custom_agent_id(self):
        ac = _make_ac()
        client = _make_mock_client()
        tracked = track_openai(client, ac, agent_id="my-bot")
        assert tracked.chat.completions._agent_id == "my-bot"


class TestTrackedCompletions:
    def test_create_calls_original(self):
        ac = _make_ac()
        response = _make_mock_response()
        client = _make_mock_client(response)
        tracked = track_openai(client, ac)

        result = tracked.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello"}],
        )

        client.chat.completions.create.assert_called_once_with(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello"}],
        )
        assert result is response

    def test_extracts_token_usage(self, capsys):
        ac = _make_ac()
        response = _make_mock_response(
            model="gpt-4o", prompt_tokens=150, completion_tokens=75
        )
        client = _make_mock_client(response)
        tracked = track_openai(client, ac)

        tracked.chat.completions.create(model="gpt-4o", messages=[])

        # Debug mode prints to stderr
        captured = capsys.readouterr()
        assert "gpt-4o" in captured.err
        assert "150" in captured.err
        assert "75" in captured.err

    def test_calculates_cost(self):
        ac = _make_ac()
        response = _make_mock_response(
            model="gpt-4o", prompt_tokens=1000, completion_tokens=500
        )
        client = _make_mock_client(response)
        tracked = track_openai(client, ac)

        with patch.object(ac._client, "enqueue") as mock_enqueue:
            tracked.chat.completions.create(model="gpt-4o", messages=[])

            event = mock_enqueue.call_args[0][0]
            assert event.tokens_in == 1000
            assert event.tokens_out == 500
            assert event.cost_usd > 0
            assert event.event_type == "llm_call"
            assert event.component_name == "gpt-4o"

    def test_handles_missing_usage(self):
        ac = _make_ac()
        response = MagicMock()
        response.model = "gpt-4o"
        response.usage = None
        client = _make_mock_client(response)
        tracked = track_openai(client, ac)

        with patch.object(ac._client, "enqueue") as mock_enqueue:
            tracked.chat.completions.create(model="gpt-4o", messages=[])

            event = mock_enqueue.call_args[0][0]
            assert event.tokens_in == 0
            assert event.tokens_out == 0

    def test_enqueues_event(self):
        ac = _make_ac()
        client = _make_mock_client()
        tracked = track_openai(client, ac)

        with patch.object(ac._client, "enqueue") as mock_enqueue:
            tracked.chat.completions.create(model="gpt-4o", messages=[])
            assert mock_enqueue.called

    def test_links_to_active_run(self):
        ac = _make_ac()
        client = _make_mock_client()
        tracked = track_openai(client, ac, agent_id="standalone")

        tracker = ac.tracker("run-agent")
        with tracker.run() as run:
            with patch.object(ac._client, "enqueue") as mock_enqueue:
                tracked.chat.completions.create(model="gpt-4o", messages=[])

                event = mock_enqueue.call_args[0][0]
                # Should use the run's agent_id, not the standalone one
                assert event.agent_id == "run-agent"
                assert event.run_id == run._run_id

    def test_measures_duration(self):
        ac = _make_ac()
        client = _make_mock_client()
        tracked = track_openai(client, ac)

        with patch.object(ac._client, "enqueue") as mock_enqueue:
            tracked.chat.completions.create(model="gpt-4o", messages=[])

            event = mock_enqueue.call_args[0][0]
            assert event.duration_ms >= 0


class TestTrackedChat:
    def test_preserves_other_chat_methods(self):
        ac = _make_ac()
        client = _make_mock_client()
        client.chat.some_other_method = MagicMock(return_value="test")
        tracked = track_openai(client, ac)
        assert tracked.chat.some_other_method() == "test"
