"""Unit tests for _sidecar.py asyncio tasks."""

import asyncio
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from tyko._sidecar import (
    SidecarConfig,
    _api_sender,
    _command_poller,
    _health_watcher,
    _system_monitor,
    stop_sentinel_path,
)


@pytest.fixture
def config() -> SidecarConfig:
    return SidecarConfig(
        client_uuid="test-client-uuid",
        api_url="https://api.tyko-labs.com",
        api_key="test-key",
        main_pid=99999,
        socket_path="/tmp/tyko-test.sock",
        system_monitor_interval=0.05,
    )


@pytest.fixture
def state() -> dict[str, Any]:
    return {"run_id": "run-123", "system_monitor_interval": 0.05}


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


def _run(coro: Any) -> Any:
    """Run a coroutine synchronously."""
    return asyncio.run(coro)


async def _drain(
    messages: list[dict[str, Any]],
    config: SidecarConfig,
    state: dict[str, Any],
) -> MagicMock:
    """Run _api_sender with a pre-filled queue and a mock HTTP client."""
    queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
    shutdown_event = asyncio.Event()

    for msg in messages:
        await queue.put(msg)

    mock_response = MagicMock(status_code=200)
    mock_client = AsyncMock()
    mock_client.post = AsyncMock(return_value=mock_response)
    mock_client.patch = AsyncMock(return_value=mock_response)
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)

    with patch("tyko._sidecar.httpx") as mock_httpx:
        mock_httpx.AsyncClient.return_value = mock_client
        task = asyncio.create_task(_api_sender(config=config, queue=queue, shutdown_event=shutdown_event, state=state))
        try:
            await asyncio.wait_for(shutdown_event.wait(), timeout=2.0)
        except asyncio.TimeoutError:
            task.cancel()

    return mock_client


# ---------------------------------------------------------------------------
# _api_sender tests
# ---------------------------------------------------------------------------


def test_api_sender_batches_log_entries(config: SidecarConfig, state: dict[str, Any]):
    """Multiple log entries within the batch window are sent as one POST."""
    messages = [
        {"type": "log", "values": {"loss": 0.5}, "timestamp": 1, "relative_time_ns": 100},
        {"type": "log", "values": {"loss": 0.4}, "timestamp": 2, "relative_time_ns": 200},
        {"type": "shutdown"},
    ]

    async def run():
        return await _drain(messages, config, state)

    client = _run(run())
    post_calls = client.post.call_args_list
    assert len(post_calls) == 1
    payload = post_calls[0].kwargs.get("json") or post_calls[0].args[1]
    assert len(payload["entries"]) == 2


def test_api_sender_patches_params(config: SidecarConfig, state: dict[str, Any]):
    messages = [
        {"type": "params", "values": {"lr": 0.001}},
        {"type": "shutdown"},
    ]
    client = _run(_drain(messages, config, state))
    patch_calls = client.patch.call_args_list
    assert any((c.kwargs.get("json") or {}).get("params") == {"lr": 0.001} for c in patch_calls)


def test_api_sender_patches_environment(config: SidecarConfig, state: dict[str, Any]):
    messages = [
        {"type": "environment", "values": {"gpu": "A100"}},
        {"type": "shutdown"},
    ]
    client = _run(_drain(messages, config, state))
    patch_calls = client.patch.call_args_list
    assert any((c.kwargs.get("json") or {}).get("environment") == {"gpu": "A100"} for c in patch_calls)


def test_api_sender_close_flushes_and_patches(config: SidecarConfig, state: dict[str, Any]):
    """close message flushes pending log batch then sends final status PATCH."""
    messages = [
        {"type": "log", "values": {"loss": 0.3}, "timestamp": 1, "relative_time_ns": 50},
        {"type": "close", "status": "completed", "at": "2026-03-15T00:00:00+00:00"},
        {"type": "shutdown"},
    ]
    client = _run(_drain(messages, config, state))
    assert client.post.called
    patch_calls = client.patch.call_args_list
    assert any((c.kwargs.get("json") or {}).get("status") == "completed" for c in patch_calls)


def test_api_sender_skips_log_when_no_run_id(config: SidecarConfig):
    """If run_id is not set, log entries are silently dropped."""
    state: dict[str, Any] = {"run_id": None}
    messages = [
        {"type": "log", "values": {"loss": 0.5}, "timestamp": 1, "relative_time_ns": 100},
        {"type": "shutdown"},
    ]
    client = _run(_drain(messages, config, state))
    assert not client.post.called


# ---------------------------------------------------------------------------
# _health_watcher tests
# ---------------------------------------------------------------------------


def test_health_watcher_enqueues_close_on_dead_pid(config: SidecarConfig):
    """If main PID is dead, health watcher enqueues close(failed) + shutdown."""

    async def run():
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        shutdown_event = asyncio.Event()

        async def fast_wait_for(coro: Any, timeout: Any = None) -> None:
            raise asyncio.TimeoutError

        with patch("tyko._sidecar.os.kill", side_effect=ProcessLookupError):
            with patch("tyko._sidecar.asyncio.wait_for", side_effect=fast_wait_for):
                await _health_watcher(config=config, queue=queue, shutdown_event=shutdown_event)

        messages = []
        while not queue.empty():
            messages.append(await queue.get())
        return messages

    messages = _run(run())
    types = [m["type"] for m in messages]
    assert "close" in types
    assert "shutdown" in types
    assert next(m for m in messages if m["type"] == "close")["status"] == "failed"


def test_health_watcher_exits_on_shutdown_event(config: SidecarConfig):
    """Health watcher exits immediately when shutdown_event is already set."""

    async def run():
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        shutdown_event = asyncio.Event()
        shutdown_event.set()
        await _health_watcher(config=config, queue=queue, shutdown_event=shutdown_event)
        return queue.empty()

    assert _run(run()) is True


# ---------------------------------------------------------------------------
# _system_monitor tests
# ---------------------------------------------------------------------------


def test_system_monitor_enqueues_log_messages(config: SidecarConfig, state: dict[str, Any]):
    """System monitor enqueues log messages with system/ prefix metrics."""
    fake_metrics = {"system/cpu_percent": 42.0, "system/ram_percent": 55.0}

    async def run():
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        shutdown_event = asyncio.Event()
        with patch("tyko._sidecar._collect_system_metrics", return_value=fake_metrics):
            task = asyncio.create_task(
                _system_monitor(config=config, queue=queue, shutdown_event=shutdown_event, state=state)
            )
            await asyncio.sleep(0.2)
            shutdown_event.set()
            await asyncio.wait_for(task, timeout=1.0)
        messages = []
        while not queue.empty():
            messages.append(await queue.get())
        return messages

    messages = _run(run())
    assert len(messages) >= 1
    assert messages[0]["type"] == "log"
    assert messages[0]["values"]["system/cpu_percent"] == 42.0


def test_system_monitor_waits_for_run_id(config: SidecarConfig):
    """System monitor does not sample until run_id is set."""

    async def run():
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        shutdown_event = asyncio.Event()
        state: dict[str, Any] = {"run_id": None, "system_monitor_interval": 0.05}
        with patch("tyko._sidecar._collect_system_metrics") as mock_collect:
            task = asyncio.create_task(
                _system_monitor(config=config, queue=queue, shutdown_event=shutdown_event, state=state)
            )
            await asyncio.sleep(0.15)
            shutdown_event.set()
            await asyncio.wait_for(task, timeout=1.0)
        return mock_collect.call_count

    assert _run(run()) == 0


def test_system_monitor_disabled_when_interval_is_none(config: SidecarConfig):
    """System monitor does not sample when interval is None."""

    async def run():
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        shutdown_event = asyncio.Event()
        state: dict[str, Any] = {"run_id": "run-123", "system_monitor_interval": None}
        with patch("tyko._sidecar._collect_system_metrics") as mock_collect:
            task = asyncio.create_task(
                _system_monitor(config=config, queue=queue, shutdown_event=shutdown_event, state=state)
            )
            await asyncio.sleep(0.1)
            shutdown_event.set()
            await asyncio.wait_for(task, timeout=1.0)
        return mock_collect.call_count

    assert _run(run()) == 0


# ---------------------------------------------------------------------------
# _command_poller tests
# ---------------------------------------------------------------------------


def test_command_poller_writes_sentinel_on_stop_command(config: SidecarConfig, tmp_path: Any):
    """Command poller writes sentinel file when stop command is received."""
    sentinel = stop_sentinel_path(client_uuid=config.client_uuid)

    # Clean up before test
    try:
        import os as _os

        _os.unlink(sentinel)
    except FileNotFoundError:
        pass

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"commands": [{"type": "stop"}]}

    mock_client = AsyncMock()
    mock_client.get = AsyncMock(return_value=mock_response)
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)

    async def run():
        shutdown_event = asyncio.Event()
        state: dict[str, Any] = {"run_id": "run-123"}
        fast_config = SidecarConfig(
            client_uuid=config.client_uuid,
            api_url=config.api_url,
            api_key=config.api_key,
            main_pid=config.main_pid,
            socket_path=config.socket_path,
            command_poll_interval=0.05,
        )
        with patch("tyko._sidecar.httpx") as mock_httpx:
            mock_httpx.AsyncClient.return_value = mock_client
            task = asyncio.create_task(_command_poller(config=fast_config, shutdown_event=shutdown_event, state=state))
            await asyncio.sleep(0.15)
            shutdown_event.set()
            await asyncio.wait_for(task, timeout=1.0)

    _run(run())

    import os as _os

    assert _os.path.exists(sentinel)

    # Clean up
    try:
        _os.unlink(sentinel)
    except FileNotFoundError:
        pass


def test_command_poller_skips_when_no_run_id(config: SidecarConfig):
    """Command poller does not poll when run_id is not set."""

    async def run():
        shutdown_event = asyncio.Event()
        state: dict[str, Any] = {"run_id": None}
        fast_config = SidecarConfig(
            client_uuid=config.client_uuid,
            api_url=config.api_url,
            api_key=config.api_key,
            main_pid=config.main_pid,
            socket_path=config.socket_path,
            command_poll_interval=0.05,
        )

        mock_client = AsyncMock()
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        with patch("tyko._sidecar.httpx") as mock_httpx:
            mock_httpx.AsyncClient.return_value = mock_client
            task = asyncio.create_task(_command_poller(config=fast_config, shutdown_event=shutdown_event, state=state))
            await asyncio.sleep(0.15)
            shutdown_event.set()
            await asyncio.wait_for(task, timeout=1.0)

        return mock_client.get.call_count

    assert _run(run()) == 0


def test_command_poller_exits_on_shutdown(config: SidecarConfig):
    """Command poller exits promptly when shutdown_event is set."""

    async def run():
        shutdown_event = asyncio.Event()
        state: dict[str, Any] = {"run_id": "run-123"}
        fast_config = SidecarConfig(
            client_uuid=config.client_uuid,
            api_url=config.api_url,
            api_key=config.api_key,
            main_pid=config.main_pid,
            socket_path=config.socket_path,
            command_poll_interval=10.0,  # Long interval — should still exit fast
        )

        mock_client = AsyncMock()
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        with patch("tyko._sidecar.httpx") as mock_httpx:
            mock_httpx.AsyncClient.return_value = mock_client
            task = asyncio.create_task(_command_poller(config=fast_config, shutdown_event=shutdown_event, state=state))
            shutdown_event.set()  # Immediately
            await asyncio.wait_for(task, timeout=1.0)

        return mock_client.get.call_count

    # Should have exited without polling
    assert _run(run()) == 0
