"""Tests for WZRD earn callback."""

from __future__ import annotations

import sys
import os
from unittest.mock import MagicMock, patch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import wzrd_earn_callback as wec


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _kwargs(model: str = "openrouter/qwen/qwen-3.5-9b", cost: float = 0.001):
    return {
        "litellm_params": {"model": model},
        "model": model,
        "response_cost": cost,
    }


def _response(ms: int = 150):
    resp = MagicMock()
    resp._response_ms = ms
    return resp


# ---------------------------------------------------------------------------
# Extraction
# ---------------------------------------------------------------------------

def test_extract_model_id_from_litellm_params():
    model = wec._extract_model_id({"litellm_params": {"model": "openrouter/qwen/qwen-3.5-9b"}})
    assert model == "openrouter/qwen/qwen-3.5-9b"


def test_extract_model_id_fallback():
    model = wec._extract_model_id({"model": "gpt-4o"})
    assert model == "gpt-4o"


def test_extract_model_id_missing():
    model = wec._extract_model_id({})
    assert model is None


def test_extract_latency_from_response():
    resp = _response(250)
    ms = wec._extract_latency_ms({}, resp)
    assert ms == 250


def test_extract_latency_from_timestamps():
    from datetime import datetime
    start = datetime(2026, 3, 23, 12, 0, 0)
    end = datetime(2026, 3, 23, 12, 0, 1, 500000)  # 1.5s
    ms = wec._extract_latency_ms({"start_time": start, "end_time": end}, MagicMock(spec=[]))
    assert ms == 1500


def test_extract_cost():
    cost = wec._extract_cost({"response_cost": 0.0032}, MagicMock())
    assert cost == 0.0032


def test_extract_cost_missing():
    cost = wec._extract_cost({}, MagicMock())
    assert cost is None


# ---------------------------------------------------------------------------
# Callback
# ---------------------------------------------------------------------------

def test_callback_disabled():
    """Disabled callback should not attempt to report."""
    os.environ["WZRD_EARN_ENABLED"] = "false"
    try:
        cb = wec.WZRDEarnCallback()
        assert not cb.enabled
        # Should return immediately without error
        cb.log_success_event(_kwargs(), _response(), None, None)
        assert cb.stats["reports_sent"] == 0
    finally:
        os.environ.pop("WZRD_EARN_ENABLED", None)


def test_callback_reports_on_success():
    """Callback should call agent.report() on success."""
    mock_agent = MagicMock()
    mock_agent.report.return_value = {"ok": True}

    # Patch _get_agent to return our mock
    with patch.object(wec, "_get_agent", return_value=mock_agent):
        cb = wec.WZRDEarnCallback(task_type="test")
        cb.log_success_event(_kwargs(), _response(200), None, None)

        mock_agent.report.assert_called_once()
        call_kwargs = mock_agent.report.call_args[1]
        assert call_kwargs["model_id"] == "openrouter/qwen/qwen-3.5-9b"
        assert call_kwargs["task_type"] == "test"
        assert call_kwargs["latency_ms"] == 200
        assert call_kwargs["cost_usd"] == 0.001
        assert call_kwargs["outcome"] == "success"
        # Structured metadata should be present
        assert call_kwargs["metadata"]["wzrd"]["source"] == "litellm-plugin"
        assert call_kwargs["metadata"]["wzrd"]["policy_version"] == "1.0"
        assert call_kwargs["metadata"]["inference"]["provider"] == "openrouter"
        assert cb.stats["reports_sent"] == 1


def test_callback_reports_failure():
    """Callback should report failures with outcome=failure."""
    mock_agent = MagicMock()
    mock_agent.report.return_value = {"ok": True}

    with patch.object(wec, "_get_agent", return_value=mock_agent):
        cb = wec.WZRDEarnCallback()
        cb.log_failure_event(_kwargs(), _response(), None, None)

        mock_agent.report.assert_called_once_with(
            model_id="openrouter/qwen/qwen-3.5-9b",
            task_type="inference",
            outcome="failure",
        )


def test_callback_survives_report_error():
    """Report failure should not crash the callback."""
    mock_agent = MagicMock()
    mock_agent.report.side_effect = Exception("network error")

    with patch.object(wec, "_get_agent", return_value=mock_agent):
        cb = wec.WZRDEarnCallback()
        # Should not raise
        cb.log_success_event(_kwargs(), _response(), None, None)
        assert cb.stats["errors"] == 1
        assert "network error" in cb.stats["last_error"]


def test_callback_skips_when_agent_init_fails():
    """If agent init failed (sentinel False), callback should no-op."""
    with patch.object(wec, "_get_agent", return_value=False):
        cb = wec.WZRDEarnCallback()
        cb.log_success_event(_kwargs(), _response(), None, None)
        assert cb.stats["reports_sent"] == 0


def test_callback_skips_missing_model():
    """No model_id → skip report."""
    mock_agent = MagicMock()
    with patch.object(wec, "_get_agent", return_value=mock_agent):
        cb = wec.WZRDEarnCallback()
        cb.log_success_event({}, _response(), None, None)
        mock_agent.report.assert_not_called()


def test_stats():
    """Stats should reflect current state."""
    cb = wec.WZRDEarnCallback()
    assert cb.stats == {
        "reports_sent": 0,
        "errors": 0,
        "last_error": None,
        "enabled": True,
    }


# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)]
    passed = failed = 0
    for test in tests:
        try:
            test()
            print(f"  PASS  {test.__name__}")
            passed += 1
        except Exception as e:
            print(f"  FAIL  {test.__name__}: {e}")
            failed += 1
    print(f"\n{passed} passed, {failed} failed out of {len(tests)}")
    sys.exit(1 if failed else 0)
