"""Tests for daemon: Unix socket server, request handler, process orchestrator."""

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

from rpipe.daemon import RPipeDaemon, DaemonError
from rpipe.protocol import (
    pack_frame,
    TYPE_RUN_REQ,
    TYPE_META,
    TYPE_STDOUT,
    TYPE_STDERR,
    TYPE_EXIT,
    TYPE_ERROR,
    TYPE_STDIN,
)
from rpipe.ssh_backend import RunHandle


class TestRPipeDaemon:
    """Test RPipeDaemon class."""

    def test_daemon_init(self):
        """Initialize daemon with socket path."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")
        assert daemon.sock_path == "/tmp/test.sock"
        assert daemon._next_run_id == 1
        assert daemon._runs == {}

    def test_alloc_run_id(self):
        """Allocate sequential run IDs."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")
        rid1 = daemon._alloc_run_id()
        rid2 = daemon._alloc_run_id()
        rid3 = daemon._alloc_run_id()

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

    @pytest.mark.asyncio
    async def test_handle_client_invalid_first_frame(self):
        """First frame must be RUN_REQ with run_id=0."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        # Create reader/writer with wrong first frame type
        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # Send META frame (wrong type) with run_id=0
        bad_frame = pack_frame(TYPE_META, 0, b"")
        reader.feed_data(bad_frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        # Should send ERROR frame
        assert writer.write.called
        error_frame = writer.write.call_args[0][0]
        # Verify it's an error frame
        assert error_frame[4] == TYPE_ERROR  # frame type at byte 4

    @pytest.mark.asyncio
    async def test_handle_client_invalid_run_id(self):
        """First frame must have run_id=0."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # Send RUN_REQ with non-zero run_id
        bad_frame = pack_frame(TYPE_RUN_REQ, 999, b"")
        reader.feed_data(bad_frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        # Should send ERROR frame
        assert writer.write.called
        error_frame = writer.write.call_args[0][0]
        assert error_frame[4] == TYPE_ERROR

    @pytest.mark.asyncio
    async def test_handle_client_invalid_json(self):
        """Invalid JSON in RUN_REQ should error."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # Send RUN_REQ with invalid JSON
        bad_json = b"not valid json"
        bad_frame = pack_frame(TYPE_RUN_REQ, 0, bad_json)
        reader.feed_data(bad_frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        assert writer.write.called
        # Verify error was sent
        error_frame = writer.write.call_args[0][0]
        assert error_frame[4] == TYPE_ERROR

    @pytest.mark.asyncio
    async def test_handle_client_missing_required_field(self):
        """Missing required field in RUN_REQ should error."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # RUN_REQ missing 'cmd' field
        req = {"host": "example.com"}
        bad_frame = pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8"))
        reader.feed_data(bad_frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        assert writer.write.called
        error_frame = writer.write.call_args[0][0]
        assert error_frame[4] == TYPE_ERROR

    @pytest.mark.asyncio
    async def test_handle_client_unsafe_cwd(self):
        """Unsafe cwd (with shell metacharacters) should be rejected."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # RUN_REQ with shell injection in cwd
        req = {
            "host": "example.com",
            "cmd": "echo hello",
            "cwd": "/tmp'; rm -rf /",  # Has quote metacharacter
        }
        frame = pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8"))
        reader.feed_data(frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        assert writer.write.called
        error_frame = writer.write.call_args[0][0]
        assert error_frame[4] == TYPE_ERROR

    @pytest.mark.asyncio
    async def test_handle_client_command_blocked_by_policy(self):
        """Commands blocked by policy should be rejected."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # RUN_REQ with rm command (denied by policy)
        req = {
            "host": "example.com",
            "cmd": "rm /tmp/file.txt",
        }
        frame = pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8"))
        reader.feed_data(frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        assert writer.write.called
        error_frame = writer.write.call_args[0][0]
        assert error_frame[4] == TYPE_ERROR

    @pytest.mark.asyncio
    async def test_handle_client_ssh_error(self):
        """SSH errors should be communicated to client."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # Mock the backend to raise SSHError
        mock_backend = AsyncMock()
        mock_backend.run.side_effect = Exception("Connection refused")
        daemon.backend = mock_backend

        # Valid RUN_REQ
        req = {
            "host": "example.com",
            "cmd": "echo hello",
        }
        frame = pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8"))
        reader.feed_data(frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        # Should send error after META
        calls = writer.write.call_args_list
        assert len(calls) >= 1
        # At least one frame should be ERROR type
        error_sent = any(call[0][0][4] == TYPE_ERROR for call in calls)
        assert error_sent

    @pytest.mark.asyncio
    async def test_handle_client_successful_run(self):
        """Successful command execution should stream output."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # Mock the backend
        mock_proc = AsyncMock()
        mock_proc.stdout = AsyncMock()
        mock_proc.stderr = AsyncMock()
        mock_proc.stdin = AsyncMock()

        # Make stdout/stderr return empty on first read (end of stream)
        async def mock_read(size):
            return b""

        mock_proc.stdout.read = mock_read
        mock_proc.stderr.read = mock_read

        mock_run = AsyncMock(spec=RunHandle)
        mock_run.proc = mock_proc
        mock_run.wait = AsyncMock(return_value=0)

        mock_backend = AsyncMock()
        mock_backend.run = AsyncMock(return_value=mock_run)
        daemon.backend = mock_backend

        # Valid RUN_REQ
        req = {
            "host": "example.com",
            "cmd": "echo hello",
        }
        frame = pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8"))
        reader.feed_data(frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        # Verify frames were sent
        assert writer.write.called

        # Should have META, EXIT frames
        calls = writer.write.call_args_list
        frame_types = [call[0][0][4] for call in calls]
        assert TYPE_META in frame_types
        assert TYPE_EXIT in frame_types

    @pytest.mark.asyncio
    async def test_handle_client_timeout_sigterm(self):
        """Process timeout should escalate from SIGTERM to SIGKILL."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        reader = asyncio.StreamReader()
        writer = AsyncMock()

        # Mock the backend with timeout
        mock_proc = AsyncMock()
        mock_proc.stdout = AsyncMock()
        mock_proc.stderr = AsyncMock()
        mock_proc.stdin = AsyncMock()

        async def mock_read(size):
            return b""

        mock_proc.stdout.read = mock_read
        mock_proc.stderr.read = mock_read

        mock_run = AsyncMock(spec=RunHandle)
        mock_run.proc = mock_proc

        # First wait_for times out, then succeeds after hard cancel
        async def mock_wait():
            await asyncio.sleep(10)
            return 0

        mock_run.wait = mock_wait
        mock_run.cancel = MagicMock()

        mock_backend = AsyncMock()
        mock_backend.run = AsyncMock(return_value=mock_run)
        daemon.backend = mock_backend

        # Valid RUN_REQ with short timeout
        req = {
            "host": "example.com",
            "cmd": "sleep 100",
            "timeout_s": 0.01,
        }
        frame = pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8"))
        reader.feed_data(frame)
        reader.feed_eof()

        await daemon.handle_client(reader, writer)

        # Verify cancel was called with hard=True after timeout
        # (This requires proper timeout handling in handle_client)
        assert mock_run.cancel.called

    @pytest.mark.asyncio
    async def test_pump_stream_writes_frames(self):
        """_pump_stream should read from stream and write frames."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        writer = AsyncMock()
        stream = AsyncMock()

        # Mock stream that returns data then EOF
        async def mock_read(size):
            if not hasattr(mock_read, "called"):
                mock_read.called = True
                return b"hello world"
            return b""

        stream.read = mock_read

        await daemon._pump_stream(writer, run_id=1, stream=stream, frame_type=TYPE_STDOUT)

        # Verify frame was written
        assert writer.write.called
        frame = writer.write.call_args[0][0]

        # Verify frame format
        assert frame[4] == TYPE_STDOUT  # frame type at byte 4
        assert frame[5:9] == struct.pack("!I", 1)  # run_id at bytes 5-8

    @pytest.mark.asyncio
    async def test_pump_stream_handles_stream_error(self):
        """_pump_stream should handle stream read errors gracefully."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        writer = AsyncMock()
        stream = AsyncMock()

        # Stream raises error on read
        stream.read = AsyncMock(side_effect=Exception("Stream error"))

        # Should not raise, just log
        await daemon._pump_stream(writer, run_id=1, stream=stream, frame_type=TYPE_STDOUT)

        # Should have attempted to write error recovery
        # (error logging is internal, just verify no exception was raised)

    @pytest.mark.asyncio
    async def test_pump_stdin_relays_stdin_frames(self):
        """_pump_stdin should relay STDIN frames to process."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        # Create reader with STDIN frame
        reader = asyncio.StreamReader()
        stdin_data = b"user input"
        stdin_frame = pack_frame(TYPE_STDIN, run_id=1, payload=stdin_data)
        reader.feed_data(stdin_frame)

        # Add a non-matching frame (wrong run_id) to ignore
        other_frame = pack_frame(TYPE_STDIN, run_id=999, payload=b"other")
        reader.feed_data(other_frame)

        # Add a second matching frame then EOF
        stdin_frame2 = pack_frame(TYPE_STDIN, run_id=1, payload=b"more")
        reader.feed_data(stdin_frame2)
        reader.feed_eof()

        # Mock run/process
        mock_proc = AsyncMock()
        mock_proc.stdin = AsyncMock()
        mock_proc.stdin.drain = AsyncMock()

        mock_run = AsyncMock(spec=RunHandle)
        mock_run.proc = mock_proc

        await daemon._pump_stdin(reader, run_id=1, run=mock_run)

        # Verify stdin received the data (should have 2 writes for matching frames)
        assert mock_proc.stdin.write.call_count == 2
        calls = [call[0][0] for call in mock_proc.stdin.write.call_args_list]
        assert calls[0] == stdin_data
        assert calls[1] == b"more"

    @pytest.mark.asyncio
    async def test_send_error(self):
        """_send_error should send ERROR frame to client."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        writer = AsyncMock()
        await daemon._send_error(writer, run_id=1, msg="test error")

        assert writer.write.called
        frame = writer.write.call_args[0][0]

        # Verify it's an ERROR frame
        assert frame[4] == TYPE_ERROR

        # Verify run_id
        (run_id,) = struct.unpack("!I", frame[5:9])
        assert run_id == 1

        # Verify payload contains error message
        payload_start = 9
        payload = frame[payload_start:]
        error_dict = json.loads(payload.decode("utf-8"))
        assert "test error" in error_dict["error"]

    @pytest.mark.asyncio
    async def test_send_error_drain_failure(self):
        """_send_error should handle drain failures gracefully."""
        daemon = RPipeDaemon(sock_path="/tmp/test.sock")

        writer = AsyncMock()
        writer.drain = AsyncMock(side_effect=Exception("Drain failed"))

        # Should not raise exception
        await daemon._send_error(writer, run_id=1, msg="test")

        # Verify frame was still written
        assert writer.write.called
