import subprocess
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from mlflow.assistant.providers.base import NotAuthenticatedError
from mlflow.assistant.providers.claude_code import ClaudeCodeProvider
from mlflow.assistant.types import EventType


class AsyncIterator:
    """Helper to mock async stdout iteration."""

    def __init__(self, items):
        self.items = iter(items)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self.items)
        except StopIteration:
            raise StopAsyncIteration


@pytest.fixture(autouse=True)
def config(tmp_path):
    config_file = tmp_path / "config.json"
    config_file.write_text('{"providers": {"claude_code": {"model": "claude-opus-4"}}}')
    with patch("mlflow.assistant.config.CONFIG_PATH", config_file):
        yield config_file


@pytest.mark.parametrize(
    ("which_return", "expected"),
    [
        ("/usr/bin/claude", True),
        (None, False),
    ],
)
def test_is_available(which_return, expected):
    with patch(
        "mlflow.assistant.providers.claude_code.shutil.which",
        return_value=which_return,
    ):
        provider = ClaudeCodeProvider()
        assert provider.is_available() is expected


def test_check_connection_runs_auth_verification_prompt():
    completed = subprocess.CompletedProcess(
        args=["claude", "-p", "hi", "--max-turns", "1", "--output-format", "json"],
        returncode=0,
        stdout='{"type": "result"}',
        stderr="",
    )

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.subprocess.run",
            return_value=completed,
        ) as mock_run,
    ):
        provider = ClaudeCodeProvider()
        provider.check_connection()

    mock_run.assert_called_once_with(
        ["claude", "-p", "hi", "--max-turns", "1", "--output-format", "json"],
        capture_output=True,
        text=True,
        timeout=30,
    )


def test_check_connection_raises_not_authenticated_for_auth_error():
    completed = subprocess.CompletedProcess(
        args=["claude", "-p", "hi", "--max-turns", "1", "--output-format", "json"],
        returncode=1,
        stdout="",
        stderr="Login required",
    )

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch("mlflow.assistant.providers.claude_code.subprocess.run", return_value=completed),
    ):
        provider = ClaudeCodeProvider()
        with pytest.raises(NotAuthenticatedError, match="claude login"):
            provider.check_connection()


@pytest.mark.asyncio
async def test_astream_yields_error_when_claude_not_found():
    with patch(
        "mlflow.assistant.providers.claude_code.shutil.which",
        return_value=None,
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    assert len(events) == 1
    assert events[0].type == EventType.ERROR
    assert "not found" in events[0].data["error"]
    assert "PATH" in events[0].data["error"]


@pytest.mark.asyncio
async def test_astream_builds_correct_command(tmp_path, monkeypatch):
    monkeypatch.setenv("TEST_ENV_VAR", "test_value")

    mock_process = MagicMock()
    mock_process.stdout = AsyncIterator([b'{"type": "result"}\n'])
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ) as mock_exec,
    ):
        provider = ClaudeCodeProvider()
        _ = [
            e async for e in provider.astream("test prompt", "http://localhost:5000", cwd=tmp_path)
        ]

    call_args = mock_exec.call_args[0]
    assert "/usr/bin/claude" in call_args
    assert "-p" in call_args
    assert "test prompt" in call_args
    assert "--output-format" in call_args
    assert "stream-json" in call_args
    assert "--verbose" in call_args
    assert "--append-system-prompt" in call_args

    # Verify system prompt contains tracking URI
    system_prompt_idx = call_args.index("--append-system-prompt") + 1
    system_prompt = call_args[system_prompt_idx]
    assert "http://localhost:5000" in system_prompt

    # Verify Skill permission is granted by default
    allowed_tools = [
        call_args[i + 1] for i, arg in enumerate(call_args) if arg == "--allowed-tools"
    ]
    assert "Skill" in allowed_tools

    # Verify cwd and tracking URI env var are passed correctly
    call_kwargs = mock_exec.call_args[1]
    assert call_kwargs["cwd"] == tmp_path
    assert call_kwargs["env"]["MLFLOW_TRACKING_URI"] == "http://localhost:5000"
    assert call_kwargs["env"]["TEST_ENV_VAR"] == "test_value"


@pytest.mark.asyncio
async def test_astream_streams_assistant_messages():
    mock_stdout = AsyncIterator([
        b'{"type": "assistant", "message": {"content": [{"type": "text", "text": "Hi!"}]}}\n',
        b'{"type": "result", "session_id": "session-123"}\n',
    ])

    mock_process = MagicMock()
    mock_process.stdout = mock_stdout
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ),
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    assert len(events) == 2
    assert events[0].type == EventType.MESSAGE
    assert events[0].data["message"]["content"][0]["text"] == "Hi!"
    assert events[1].type == EventType.DONE
    assert events[1].data["session_id"] == "session-123"


@pytest.mark.asyncio
async def test_astream_emits_usage_event_before_done():
    mock_stdout = AsyncIterator([
        b'{"type": "assistant", "message": {"content": [{"type": "text", "text": "Hi!"}]}}\n',
        b'{"type": "result", "session_id": "session-123", "total_cost_usd": 0.1319, '
        b'"usage": {"input_tokens": 2, "cache_creation_input_tokens": 35155, '
        b'"cache_read_input_tokens": 100, "output_tokens": 5}}\n',
    ])

    mock_process = MagicMock()
    mock_process.stdout = mock_stdout
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ),
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    usage_events = [
        e
        for e in events
        if e.type == EventType.STREAM_EVENT and e.data["event"].get("type") == "usage"
    ]
    assert len(usage_events) == 1
    assert usage_events[0].data["event"]["usage"] == {
        "prompt_tokens": 35257,
        "completion_tokens": 5,
        "total_tokens": 35262,
        "total_cost_usd": 0.1319,
    }
    # The usage event must precede the DONE event, which closes the client stream.
    assert events.index(usage_events[0]) < next(
        i for i, e in enumerate(events) if e.type == EventType.DONE
    )


@pytest.mark.parametrize(
    ("usage", "cost_usd", "expected"),
    [
        (
            {"input_tokens": 10, "output_tokens": 5},
            0.25,
            {
                "prompt_tokens": 10,
                "completion_tokens": 5,
                "total_tokens": 15,
                "total_cost_usd": 0.25,
            },
        ),
        (
            {
                "input_tokens": 2,
                "cache_creation_input_tokens": 100,
                "cache_read_input_tokens": 50,
                "output_tokens": 5,
            },
            None,
            {
                "prompt_tokens": 152,
                "completion_tokens": 5,
                "total_tokens": 157,
                "total_cost_usd": None,
            },
        ),
        (
            {},
            None,
            {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "total_cost_usd": None,
            },
        ),
    ],
)
def test_build_usage_event_sums_cache_tokens(usage, cost_usd, expected):
    event = ClaudeCodeProvider._build_usage_event(usage, cost_usd)
    assert event.type == EventType.STREAM_EVENT
    assert event.data["event"]["type"] == "usage"
    assert event.data["event"]["usage"] == expected


@pytest.mark.asyncio
async def test_astream_handles_process_error():
    mock_process = MagicMock()
    mock_process.stdout = AsyncIterator([])
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"Command failed")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 1

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ),
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    assert events[-1].type == EventType.ERROR
    assert "Command failed" in events[-1].data["error"]


@pytest.mark.asyncio
async def test_astream_passes_session_id_for_resume():
    mock_process = MagicMock()
    mock_process.stdout = AsyncIterator([b'{"type": "result"}\n'])
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ) as mock_exec,
    ):
        provider = ClaudeCodeProvider()
        _ = [
            e
            async for e in provider.astream(
                "prompt", "http://localhost:5000", session_id="existing-session"
            )
        ]

    call_args = mock_exec.call_args[0]
    assert "--resume" in call_args
    assert "existing-session" in call_args


@pytest.mark.asyncio
async def test_astream_handles_non_json_output():
    mock_stdout = AsyncIterator([
        b"Some plain text output\n",
        b'{"type": "result"}\n',
    ])

    mock_process = MagicMock()
    mock_process.stdout = mock_stdout
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ),
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    assert events[0].type == EventType.MESSAGE
    assert events[0].data["message"]["content"] == "Some plain text output"


@pytest.mark.asyncio
async def test_astream_handles_error_message_type():
    mock_stdout = AsyncIterator([
        b'{"type": "error", "error": {"message": "API rate limit exceeded"}}\n',
    ])

    mock_process = MagicMock()
    mock_process.stdout = mock_stdout
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ),
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    assert events[0].type == EventType.ERROR
    assert "rate limit" in events[0].data["error"]


@pytest.mark.asyncio
async def test_astream_skips_rate_limit_event():
    mock_stdout = AsyncIterator([
        b'{"type": "rate_limit_event", "data": {"retry_after": 1.0}}\n',
        b'{"type": "assistant", "message": {"content": [{"type": "text", "text": "hello"}]}}\n',
        b'{"type": "result", "result": null, "session_id": "sess-123"}\n',
    ])

    mock_process = MagicMock()
    mock_process.stdout = mock_stdout
    mock_process.stderr = MagicMock()
    mock_process.stderr.read = AsyncMock(return_value=b"")
    mock_process.wait = AsyncMock()
    mock_process.returncode = 0
    mock_process.pid = 12345

    with (
        patch(
            "mlflow.assistant.providers.claude_code.shutil.which",
            return_value="/usr/bin/claude",
        ),
        patch(
            "mlflow.assistant.providers.claude_code.asyncio.create_subprocess_exec",
            return_value=mock_process,
        ),
    ):
        provider = ClaudeCodeProvider()
        events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]

    assert len(events) == 2
    assert events[0].type == EventType.MESSAGE
    assert events[1].type == EventType.DONE