"""Tests for client session supporting multiple concurrent runs."""

import asyncio
import json
import struct
from unittest.mock import AsyncMock, MagicMock

import pytest

from rpipe.client_session import ClientSession, RunResult
from rpipe.protocol import pack_frame, TYPE_META, TYPE_STDOUT, TYPE_STDERR, TYPE_EXIT, TYPE_ERROR


class TestClientSessionBasic:
    """Test basic client session functionality."""

    def test_create_session(self):
        """Create a client session."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        assert session.default_host == "example.com"
        assert len(session._runs) == 0

    def test_alloc_run_id(self):
        """Allocate run IDs sequentially."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer)

        rid1 = session._alloc_run_id()
        rid2 = session._alloc_run_id()
        rid3 = session._alloc_run_id()

        assert rid1 == 1
        assert rid2 == 2
        assert rid3 == 3


@pytest.mark.asyncio
class TestClientSessionCommands:
    """Test submitting commands via session."""

    async def test_run_command_with_default_host(self):
        """Submit command using default host."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        run_id = await session.run_command("cat /var/log/syslog")

        assert run_id == 1
        assert run_id in session._runs
        assert session._runs[run_id].run_id == 1

    async def test_run_command_with_explicit_host(self):
        """Submit command with explicit host."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="default.com")

        run_id = await session.run_command(
            "grep error /var/log/syslog",
            host="custom.com",
        )

        assert run_id == 1
        assert session._run_metadata[run_id]["host"] == "custom.com"

    async def test_run_command_without_host_fails(self):
        """Submitting command without host fails if no default."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer)  # No default host

        with pytest.raises(ValueError):
            await session.run_command("cat /var/log/syslog")

    async def test_run_command_sends_frame(self):
        """Submitting command sends RUN_REQ frame."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cat /var/log/syslog", timeout_s=30, policy="read_only")

        # Should have called writer.write
        assert writer.write.called
        assert writer.drain.called

    async def test_multiple_commands_allocate_different_ids(self):
        """Multiple commands get different run IDs."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        rid1 = await session.run_command("cat /var/log/syslog")
        rid2 = await session.run_command("grep error /var/log/auth.log")
        rid3 = await session.run_command("ls /tmp")

        assert rid1 == 1
        assert rid2 == 2
        assert rid3 == 3
        assert len(session._runs) == 3


@pytest.mark.asyncio
class TestClientSessionOutput:
    """Test output handling."""

    async def test_set_stdout_callback(self):
        """Set callback for stdout."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cat /var/log/syslog")

        callback = AsyncMock()
        session.set_output_callback(1, on_stdout=callback)

        assert session._output_callbacks[1]["stdout"] == callback

    async def test_pump_responses_stdout(self):
        """Process stdout frames."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cat /var/log/syslog")

        # Queue responses
        meta_frame = pack_frame(TYPE_META, 1, json.dumps({"run_id": 1}).encode())
        stdout_frame = pack_frame(TYPE_STDOUT, 1, b"output line 1\n")
        exit_frame = pack_frame(TYPE_EXIT, 1, struct.pack("!i", 0))

        reader.readexactly.side_effect = [
            b"\x00\x00\x00\x14",  # meta length
            meta_frame[4:],  # meta body
            b"\x00\x00\x00\x1f",  # stdout length
            stdout_frame[4:],  # stdout body
            b"\x00\x00\x00\x09",  # exit length
            exit_frame[4:],  # exit body
            asyncio.CancelledError(),  # Stop pumping
        ]

        try:
            await session.pump_responses()
        except asyncio.CancelledError:
            pass

        result = session._runs[1]
        assert result.stdout == b"output line 1\n"
        assert result.exit_code == 0
        assert result.completed

    async def test_pump_responses_stderr(self):
        """Process stderr frames."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cat /var/log/syslog")

        # Queue responses
        stderr_frame = pack_frame(TYPE_STDERR, 1, b"error message\n")
        exit_frame = pack_frame(TYPE_EXIT, 1, struct.pack("!i", 1))

        reader.readexactly.side_effect = [
            b"\x00\x00\x00\x1f",  # stderr length
            stderr_frame[4:],  # stderr body
            b"\x00\x00\x00\x09",  # exit length
            exit_frame[4:],  # exit body
            asyncio.CancelledError(),
        ]

        try:
            await session.pump_responses()
        except asyncio.CancelledError:
            pass

        result = session._runs[1]
        assert result.stderr == b"error message\n"
        assert result.exit_code == 1

    async def test_pump_responses_error_frame(self):
        """Process error frames."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cat /var/log/syslog")

        # Queue error response
        error_frame = pack_frame(TYPE_ERROR, 1, json.dumps({"error": "Command rejected"}).encode())

        reader.readexactly.side_effect = [
            b"\x00\x00\x00\x25",  # error length
            error_frame[4:],  # error body
            asyncio.CancelledError(),
        ]

        try:
            await session.pump_responses()
        except asyncio.CancelledError:
            pass

        result = session._runs[1]
        assert result.error == "Command rejected"
        assert result.completed


@pytest.mark.asyncio
class TestClientSessionWaiting:
    """Test waiting for command completion."""

    async def test_wait_for_run_success(self):
        """Wait for a run to complete successfully."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("echo hello")

        # Manually mark as completed
        session._runs[1].exit_code = 0
        session._runs[1].stdout = b"hello\n"
        session._runs[1].completed = True

        result = await session.wait_for_run(1, timeout=1.0)

        assert result.run_id == 1
        assert result.exit_code == 0
        assert result.stdout == b"hello\n"

    async def test_wait_for_run_timeout(self):
        """Waiting with timeout fails if run doesn't complete."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("sleep 100")

        with pytest.raises(TimeoutError):
            await session.wait_for_run(1, timeout=0.1)

    async def test_wait_for_unknown_run(self):
        """Waiting for unknown run fails."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        with pytest.raises(KeyError):
            await session.wait_for_run(999)


@pytest.mark.asyncio
class TestClientSessionControl:
    """Test control operations (send stdin, cancel, etc)."""

    async def test_send_stdin(self):
        """Send stdin to a running command."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cat", stdin_mode="pipe")
        await session.send_stdin(1, b"hello\n")

        assert writer.write.called
        assert writer.drain.called

    async def test_cancel_command(self):
        """Cancel a running command."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("sleep 100")
        await session.cancel_command(1)

        assert writer.write.called
        assert writer.drain.called


class TestClientSessionState:
    """Test session state management."""

    def test_get_active_runs(self):
        """Get list of active runs."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        # Manually create some runs
        session._next_run_id = 1
        session._runs[1] = RunResult(run_id=1, completed=False)
        session._runs[2] = RunResult(run_id=2, completed=True)
        session._runs[3] = RunResult(run_id=3, completed=False)

        active = session.get_active_runs()
        assert set(active) == {1, 3}

    def test_get_result_not_completed(self):
        """Getting result of non-completed run fails."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        session._next_run_id = 1
        session._runs[1] = RunResult(run_id=1, completed=False)

        with pytest.raises(RuntimeError):
            session.get_result(1)

    def test_get_result_completed(self):
        """Getting result of completed run succeeds."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        session._next_run_id = 1
        session._runs[1] = RunResult(
            run_id=1, exit_code=0, stdout=b"output\n", completed=True
        )

        result = session.get_result(1)
        assert result.exit_code == 0
        assert result.stdout == b"output\n"


@pytest.mark.asyncio
class TestClientSessionConcurrency:
    """Test concurrent command execution."""

    async def test_multiple_concurrent_commands(self):
        """Submit multiple commands concurrently."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        # Submit multiple commands
        ids = await asyncio.gather(
            session.run_command("cat /var/log/syslog"),
            session.run_command("grep error /var/log/auth.log"),
            session.run_command("ls /tmp"),
        )

        assert ids == [1, 2, 3]
        assert len(session._runs) == 3

        # All should be pending
        assert len(session.get_active_runs()) == 3

    async def test_concurrent_output_processing(self):
        """Process output from multiple concurrent commands."""
        reader = AsyncMock()
        writer = AsyncMock()
        session = ClientSession(reader, writer, default_host="example.com")

        await session.run_command("cmd1")
        await session.run_command("cmd2")

        # Simulate receiving output for both
        # Run 1: stdout
        session._runs[1].stdout = b"output1\n"
        # Run 2: stdout
        session._runs[2].stdout = b"output2\n"

        # Both should be independent
        assert session._runs[1].stdout == b"output1\n"
        assert session._runs[2].stdout == b"output2\n"
