"""
Unit tests for wrap_claude_agent (Claude Agent SDK wrapper).

Tests the Python wrapper that auto-tracks sessions, tool calls,
tokens, and costs to Sentrial via message-stream observation.
"""

import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest

from sentrial.claude_code import wrap_claude_agent


# ===========================================================================
# Test data
# ===========================================================================

INIT_MESSAGE = {
    "type": "system",
    "subtype": "init",
    "session_id": "sdk-session-abc",
    "model": "claude-sonnet-4-20250514",
    "tools": ["Bash", "Read", "Write"],
    "cwd": "/home/user/project",
    "mcp_servers": ["filesystem"],
}

ASSISTANT_MESSAGE = {
    "type": "assistant",
    "message": {"role": "assistant", "content": "I will help you with that."},
}

ASSISTANT_TOOL_USE_MESSAGE = {
    "type": "assistant",
    "message": {
        "role": "assistant",
        "content": [
            {"type": "tool_use", "name": "Bash", "input": {"command": "ls"}, "id": "tu-1"},
        ],
    },
}

ASSISTANT_MULTI_TOOL_USE_MESSAGE = {
    "type": "assistant",
    "message": {
        "role": "assistant",
        "content": [
            {"type": "tool_use", "name": "Bash", "input": {"command": "ls"}, "id": "tu-2"},
            {"type": "tool_use", "name": "Read", "input": {"file_path": "/x"}, "id": "tu-3"},
        ],
    },
}

USER_TOOL_RESULT_MESSAGE = {
    "type": "user",
    "message": {
        "role": "user",
        "content": [
            {"type": "tool_result", "tool_use_id": "tu-1", "content": "file.txt"},
        ],
    },
}

USER_MULTI_TOOL_RESULT_MESSAGE = {
    "type": "user",
    "message": {
        "role": "user",
        "content": [
            {"type": "tool_result", "tool_use_id": "tu-2", "content": "file.txt\nREADME.md"},
            {"type": "tool_result", "tool_use_id": "tu-3", "content": "hello world"},
        ],
    },
}

RESULT_MESSAGE = {
    "type": "result",
    "subtype": "success",
    "is_error": False,
    "result": "Task completed successfully.",
    "total_cost_usd": 0.0234,
    "duration_ms": 5200,
    "duration_api_ms": 4800,
    "num_turns": 3,
    "usage": {"input_tokens": 1500, "output_tokens": 800},
}

ERROR_RESULT_MESSAGE = {
    "type": "result",
    "subtype": "error_during_execution",
    "is_error": True,
    "result": "",
    "total_cost_usd": 0.005,
    "duration_ms": 1200,
    "duration_api_ms": 1000,
    "num_turns": 1,
    "usage": {"input_tokens": 500, "output_tokens": 100},
    "errors": ["Tool execution failed", "Permission denied"],
}


# ===========================================================================
# Helpers
# ===========================================================================


def make_mock_query(messages):
    """Create a mock query function that yields given messages."""
    captured = {"kwargs": None}

    async def mock_query(**kwargs):
        captured["kwargs"] = kwargs
        for msg in messages:
            yield msg

    return mock_query, captured


def make_mock_client():
    """Create a mock AsyncSentrialClient."""
    client = MagicMock()
    client.create_session = AsyncMock(return_value="session-123")
    client.complete_session = AsyncMock(return_value={})
    client.track_tool_call = AsyncMock(return_value={})
    return client


async def collect(gen):
    """Collect all items from an async iterator."""
    results = []
    async for item in gen:
        results.append(item)
    return results


# ===========================================================================
# Tests
# ===========================================================================


@pytest.mark.asyncio
async def test_yields_all_messages_unchanged():
    """All messages pass through the wrapper untouched."""
    messages = [INIT_MESSAGE, ASSISTANT_MESSAGE, RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="test-agent", user_id="u1")
    yielded = await collect(wrapped(prompt="Hello"))

    assert len(yielded) == 3
    assert yielded[0] == INIT_MESSAGE
    assert yielded[1] == ASSISTANT_MESSAGE
    assert yielded[2] == RESULT_MESSAGE


@pytest.mark.asyncio
async def test_creates_session_on_init():
    """Session is created when the init message is observed."""
    messages = [INIT_MESSAGE, RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(
        query_fn, client=client, default_agent="my-agent", user_id="user-42"
    )
    await collect(wrapped(prompt="Fix the bug"))

    client.create_session.assert_called_once()
    call_kwargs = client.create_session.call_args[1]
    assert call_kwargs["agent_name"] == "my-agent"
    assert call_kwargs["user_id"] == "user-42"
    assert "Fix the bug" in call_kwargs["name"]
    assert call_kwargs["metadata"]["model"] == "claude-sonnet-4-20250514"
    assert call_kwargs["metadata"]["tools"] == ["Bash", "Read", "Write"]
    assert call_kwargs["metadata"]["sdk_session_id"] == "sdk-session-abc"


@pytest.mark.asyncio
async def test_completes_session_on_result():
    """Session is completed with metrics when the result message is observed."""
    messages = [INIT_MESSAGE, RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Do something"))

    client.complete_session.assert_called_once()
    kw = client.complete_session.call_args[1]
    assert kw["session_id"] == "session-123"
    assert kw["success"] is True
    assert kw["estimated_cost"] == 0.0234
    assert kw["prompt_tokens"] == 1500
    assert kw["completion_tokens"] == 800
    assert kw["total_tokens"] == 2300
    assert kw["duration_ms"] == 5200
    assert kw["assistant_output"] == "Task completed successfully."
    assert kw["custom_metrics"]["num_turns"] == 3
    assert kw["custom_metrics"]["duration_api_ms"] == 4800


@pytest.mark.asyncio
async def test_marks_session_failed_on_error_result():
    """Error result marks session as failed with failure_reason."""
    messages = [INIT_MESSAGE, ERROR_RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Do something"))

    client.complete_session.assert_called_once()
    kw = client.complete_session.call_args[1]
    assert kw["success"] is False
    assert "Tool execution failed" in kw["failure_reason"]
    assert "Permission denied" in kw["failure_reason"]


@pytest.mark.asyncio
async def test_tracks_tool_calls_from_message_stream():
    """Tool calls are tracked by observing AssistantMessage + UserMessage pairs."""
    messages = [
        INIT_MESSAGE,
        ASSISTANT_TOOL_USE_MESSAGE,
        USER_TOOL_RESULT_MESSAGE,
        RESULT_MESSAGE,
    ]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Hello"))

    # Let fire-and-forget tasks complete
    await asyncio.sleep(0.05)

    client.track_tool_call.assert_called_once()
    kw = client.track_tool_call.call_args[1]
    assert kw["session_id"] == "session-123"
    assert kw["tool_name"] == "Bash"
    assert kw["tool_input"] == {"command": "ls"}
    assert kw["tool_output"] == {"response": "file.txt"}
    assert kw["metadata"] == {"tool_use_id": "tu-1"}


@pytest.mark.asyncio
async def test_tracks_multiple_tool_calls_in_one_turn():
    """Multiple tool uses in a single AssistantMessage are all tracked."""
    messages = [
        INIT_MESSAGE,
        ASSISTANT_MULTI_TOOL_USE_MESSAGE,
        USER_MULTI_TOOL_RESULT_MESSAGE,
        RESULT_MESSAGE,
    ]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Hello"))

    await asyncio.sleep(0.05)

    assert client.track_tool_call.call_count == 2
    calls = [c[1] for c in client.track_tool_call.call_args_list]
    tool_names = {c["tool_name"] for c in calls}
    assert tool_names == {"Bash", "Read"}


@pytest.mark.asyncio
async def test_degrades_when_create_session_returns_none():
    """When create_session returns None, messages still pass through."""
    messages = [INIT_MESSAGE, ASSISTANT_MESSAGE, RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()
    client.create_session.return_value = None

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    yielded = await collect(wrapped(prompt="Hello"))

    assert len(yielded) == 3
    client.create_session.assert_called_once()
    client.complete_session.assert_not_called()


@pytest.mark.asyncio
async def test_degrades_when_create_session_throws():
    """When create_session raises, messages still pass through."""
    messages = [INIT_MESSAGE, ASSISTANT_MESSAGE, RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()
    client.create_session.side_effect = Exception("network error")

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    yielded = await collect(wrapped(prompt="Hello"))

    assert len(yielded) == 3
    client.create_session.assert_called_once()
    client.complete_session.assert_not_called()


@pytest.mark.asyncio
async def test_passes_convo_id_and_extra_metadata():
    """convo_id and extra_metadata are forwarded to create_session."""
    messages = [INIT_MESSAGE, RESULT_MESSAGE]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(
        query_fn,
        client=client,
        default_agent="t",
        user_id="u",
        convo_id="convo-abc",
        extra_metadata={"environment": "staging", "version": "1.2.3"},
    )
    await collect(wrapped(prompt="Hello"))

    client.create_session.assert_called_once()
    kw = client.create_session.call_args[1]
    assert kw["convo_id"] == "convo-abc"
    assert kw["metadata"]["environment"] == "staging"
    assert kw["metadata"]["version"] == "1.2.3"


@pytest.mark.asyncio
async def test_completes_as_failed_when_generator_throws():
    """If the upstream generator raises, session is marked as failed."""
    client = make_mock_client()

    async def throwing_query(**kwargs):
        yield INIT_MESSAGE
        raise RuntimeError("upstream crash")

    wrapped = wrap_claude_agent(throwing_query, client=client, default_agent="t", user_id="u")
    gen = wrapped(prompt="Hello")

    # Consume init message
    msg = await gen.__anext__()
    assert msg["type"] == "system"

    # Next should raise
    with pytest.raises(RuntimeError, match="upstream crash"):
        await gen.__anext__()

    client.complete_session.assert_called_once()
    kw = client.complete_session.call_args[1]
    assert kw["success"] is False
    assert "upstream crash" in kw["failure_reason"]


@pytest.mark.asyncio
async def test_structured_tool_responses_pass_through():
    """Dict tool results are passed through directly, not wrapped."""
    structured_result = {"files": ["a.ts", "b.ts"], "count": 2}
    messages = [
        INIT_MESSAGE,
        {
            "type": "assistant",
            "message": {
                "content": [
                    {"type": "tool_use", "name": "Glob", "input": {"pattern": "*.ts"}, "id": "tu-s1"},
                ],
            },
        },
        {
            "type": "user",
            "message": {
                "content": [
                    {"type": "tool_result", "tool_use_id": "tu-s1", "content": structured_result},
                ],
            },
        },
        RESULT_MESSAGE,
    ]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Hello"))

    await asyncio.sleep(0.05)

    kw = client.track_tool_call.call_args[1]
    assert kw["tool_output"] == {"files": ["a.ts", "b.ts"], "count": 2}


@pytest.mark.asyncio
async def test_options_passed_through_unchanged():
    """Options are forwarded to query_fn without modification."""
    messages = [INIT_MESSAGE, RESULT_MESSAGE]
    query_fn, captured = make_mock_query(messages)
    client = make_mock_client()

    my_options = {"model": "claude-haiku", "max_turns": 5}
    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Hello", options=my_options))

    assert captured["kwargs"]["options"] is my_options


@pytest.mark.asyncio
async def test_no_tool_tracking_without_session():
    """Tool calls in the stream are ignored if no session was created."""
    messages = [
        INIT_MESSAGE,
        ASSISTANT_TOOL_USE_MESSAGE,
        USER_TOOL_RESULT_MESSAGE,
        RESULT_MESSAGE,
    ]
    query_fn, _ = make_mock_query(messages)
    client = make_mock_client()
    client.create_session.return_value = None  # No session

    wrapped = wrap_claude_agent(query_fn, client=client, default_agent="t", user_id="u")
    await collect(wrapped(prompt="Hello"))

    await asyncio.sleep(0.05)

    client.track_tool_call.assert_not_called()
