"""
Unit tests for LatchContext and latch() context manager
"""

import pytest
from unittest.mock import AsyncMock, patch, MagicMock
import httpx

from latch.context import latch, LatchContext
from latch.models import DecisionResponse, DecisionMeta


@pytest.fixture
def mock_decision_response():
    return {
        "data": {
            "decisionId": "dec_test123",
            "shouldCallLLM": True,
            "model": "gpt-4o-mini",
            "contextFraction": 1.0,
            "estimatedCost": 0.0025,
            "reasoning": "Test decision",
            "metadata": {
                "inputTokens": 100,
                "outputTokensEstimate": 50,
                "wasCached": False,
                "decisionLatencyMs": 15,
            },
        }
    }


@pytest.fixture
def mock_report_response():
    return {
        "data": {
            "decisionId": "dec_test123",
            "actualCost": 0.003,
            "totalLatencyMs": 250,
        }
    }


def test_latch_context_creation():
    """Test LatchContext can be created with required params"""
    ctx = LatchContext(
        api_key="test_key",
        base_url="http://localhost:8081",
        task="qa",
    )
    assert ctx.api_key == "test_key"
    assert ctx.task == "qa"
    assert ctx.decision is None


def test_latch_context_properties_before_enter():
    """Test that properties raise before context is entered"""
    ctx = LatchContext(api_key="test_key")

    with pytest.raises(RuntimeError):
        _ = ctx.model_name

    with pytest.raises(RuntimeError):
        _ = ctx.decision_id

    with pytest.raises(RuntimeError):
        _ = ctx.should_call_llm


def test_latch_function_requires_api_key(monkeypatch):
    """Test latch() raises when no API key provided"""
    monkeypatch.delenv("LATCH_API_KEY", raising=False)
    with pytest.raises(ValueError, match="api_key is required"):
        latch(task="test")


def test_latch_function_uses_env_var(monkeypatch):
    """Test latch() uses LATCH_API_KEY env var"""
    monkeypatch.setenv("LATCH_API_KEY", "env_test_key")

    ctx = latch(task="test")
    assert ctx.api_key == "env_test_key"


def test_latch_context_with_all_params():
    """Test LatchContext with all optional params"""
    ctx = LatchContext(
        api_key="test_key",
        base_url="http://custom:9000",
        task="code-generation",
        budget=0.50,
        max_latency=3000,
        metadata={"user": "test"},
        timeout=60.0,
    )
    assert ctx.base_url == "http://custom:9000"
    assert ctx.task == "code-generation"
    assert ctx.budget == 0.50
    assert ctx.max_latency == 3000
    assert ctx.metadata == {"user": "test"}
    assert ctx.timeout == 60.0


def test_record_usage():
    """Test record_usage stores data correctly"""
    ctx = LatchContext(api_key="test_key")
    ctx.record_usage(input_tokens=100, output_tokens=50, model="custom-model")

    assert ctx._usage["input_tokens"] == 100
    assert ctx._usage["output_tokens"] == 50
    assert ctx._usage["model"] == "custom-model"


@pytest.mark.asyncio
async def test_context_manager_flow(mock_decision_response, mock_report_response):
    """Test full context manager enter/exit flow"""

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json.return_value = mock_decision_response

    mock_report_resp = MagicMock()
    mock_report_resp.status_code = 200
    mock_report_resp.json.return_value = mock_report_response

    with patch("httpx.AsyncClient") as mock_client_class:
        mock_client = AsyncMock()
        mock_client.post = AsyncMock(side_effect=[mock_response, mock_report_resp])
        mock_client.aclose = AsyncMock()
        mock_client_class.return_value = mock_client

        async with latch(api_key="test_key", task="qa") as ctx:
            assert ctx.decision_id == "dec_test123"
            assert ctx.model_name == "gpt-4o-mini"
            assert ctx.should_call_llm is True
            assert ctx.estimated_cost == 0.0025

            ctx.record_usage(input_tokens=100, output_tokens=50)

        assert mock_client.post.call_count == 2

        decision_call = mock_client.post.call_args_list[0]
        assert decision_call[0][0] == "/v1/decision"

        report_call = mock_client.post.call_args_list[1]
        assert report_call[0][0] == "/v1/report"
        report_body = report_call[1]["json"]
        assert report_body["decisionId"] == "dec_test123"
        assert report_body["inputTokens"] == 100
        assert report_body["outputTokens"] == 50


@pytest.mark.asyncio
async def test_context_manager_handles_error(mock_decision_response):
    """Test context manager still reports on error"""

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json.return_value = mock_decision_response

    mock_report_resp = MagicMock()
    mock_report_resp.status_code = 200

    with patch("httpx.AsyncClient") as mock_client_class:
        mock_client = AsyncMock()
        mock_client.post = AsyncMock(side_effect=[mock_response, mock_report_resp])
        mock_client.aclose = AsyncMock()
        mock_client_class.return_value = mock_client

        try:
            async with latch(api_key="test_key", task="error-test") as ctx:
                ctx.record_usage(input_tokens=50, output_tokens=0)
                raise ValueError("Test error")
        except ValueError:
            pass

        assert mock_client.post.call_count == 2
        report_call = mock_client.post.call_args_list[1]
        report_body = report_call[1]["json"]
        assert report_body["success"] is False
        assert "Test error" in report_body["errorMessage"]
