"""Module to test the Timepix Fly client functionality."""

from __future__ import annotations

from unittest import mock

import pytest
from ophyd import StatusBase
from websockets import State

from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client import (
    TimepixFlyClient,
    TimePixFlyStatus,
    TimePixStatusError,
)
from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import ProgramState


@pytest.fixture(scope="function")
def timepix_fly_client():
    """Fixture for creating a Timepix Fly client instance."""
    client = TimepixFlyClient(rest_url="http://localhost:8000", ws_url="ws://localhost:8000/ws")
    try:
        yield client
    finally:
        client.shutdown()


@pytest.mark.parametrize(
    "return_state",
    [
        ProgramState(state=TimePixFlyStatus.CONFIG),
        ProgramState(state=TimePixFlyStatus.COLLECT),
        ProgramState(state=TimePixFlyStatus.SETUP),
    ],
)
def test_timepix_fly_client_stop_running_collection(timepix_fly_client, return_state):
    """Test the on_connected method of the Timepix Fly client."""
    with (
        mock.patch.object(timepix_fly_client, "stop_collect") as mock_stop_collect,
        mock.patch.object(timepix_fly_client, "state", return_value=return_state),
    ):
        timepix_fly_client.stop_running_collection()

        if return_state.state == TimePixFlyStatus.CONFIG:
            assert mock_stop_collect.call_count == 0, "Stop collect should be called once."
            timepix_fly_client._started = True
            timepix_fly_client.stop_running_collection()
            assert mock_stop_collect.call_count == 1, "Stop collect should not be called again."
        else:
            assert mock_stop_collect.call_count == 1, "Stop collect should be called once."
            timepix_fly_client._started = True
            timepix_fly_client.stop_running_collection()
            assert mock_stop_collect.call_count == 2, "Stop collect should be called in CONFIG."


def test_timepix_fly_client_on_connected(timepix_fly_client):
    """
    Test timepix fly client connect method.

    This simply ensures that all methods are called. They are tested separately.
    """
    with (
        mock.patch.object(timepix_fly_client, "stop_running_collection") as mock_stop_collection,
        mock.patch.object(timepix_fly_client, "connect") as mock_connect,
        mock.patch.object(timepix_fly_client, "wait_for_connection") as mock_wait_for_connection,
    ):
        timepix_fly_client.on_connected()
        mock_stop_collection.assert_called_once()
        mock_connect.assert_called_once()
        mock_wait_for_connection.assert_called_once()


def test_timepix_fly_client_connect(timepix_fly_client):
    """This tests the connect method of timepix fly client."""
    module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
    # Patch Thread so we don't actually start the background loop
    with mock.patch(f"{module_path}.threading.Thread") as mock_thread_cls:
        mock_thread = mock.Mock()
        mock_thread_cls.return_value = mock_thread

        timepix_fly_client.connect()

        # Thread should be created with the update loop as target and daemon True
        mock_thread_cls.assert_called_once()
        # start() must be called on the created thread
        mock_thread.start.assert_called_once()


def test_timepix_fly_client_wait_for_connection(timepix_fly_client):
    """This tests the wait_for_connection method of timepix fly client."""
    module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
    # Case 1: ws_client already present and OPEN
    mock_client = mock.Mock()
    type(mock_client).state = mock.PropertyMock(return_value=State.OPEN)
    timepix_fly_client.ws_client = mock_client

    # Should return immediately
    timepix_fly_client.wait_for_connection(timeout=0.1)

    # Case 2: connect() establishes the connection
    timepix_fly_client.ws_client = None
    mock_client2 = mock.Mock()
    type(mock_client2).state = mock.PropertyMock(return_value=State.OPEN)
    with mock.patch(f"{module_path}.connect", return_value=mock_client2) as mock_connect:
        timepix_fly_client.wait_for_connection(timeout=0.1)
        mock_connect.assert_called_once()


def test_timepix_fly_client_ws_send_and_received(timepix_fly_client):
    """This tests the _ws_send_and_receive method of timepix fly client."""
    # Prepare a mock ws_client where recv returns a message
    mock_ws = mock.Mock()
    mock_ws.recv.return_value = "init"
    timepix_fly_client.ws_client = mock_ws

    # Patch _on_received_ws_message to ensure it's invoked
    with mock.patch.object(timepix_fly_client, "_on_received_ws_message") as mock_on_msg:
        timepix_fly_client._ws_send_and_receive()
        mock_on_msg.assert_called_once_with("init")

    # Now simulate a TimeoutError from recv: should simply return and not call _on_received_ws_message
    mock_ws2 = mock.Mock()
    mock_ws2.recv.side_effect = TimeoutError
    timepix_fly_client.ws_client = mock_ws2
    with mock.patch.object(timepix_fly_client, "_on_received_ws_message") as mock_on_msg2:
        timepix_fly_client._ws_send_and_receive()
        mock_on_msg2.assert_not_called()


def test_timepix_fly_client_on_message_received(timepix_fly_client):
    """This tests the _on_received_ws_message method of timepix fly client."""
    with mock.patch.object(timepix_fly_client, "_run_status_callbacks") as mock_run_callbacks:
        timepix_fly_client._on_received_ws_message("init")
        assert timepix_fly_client._status == TimePixFlyStatus.INIT
        mock_run_callbacks.assert_called_once()

    # invalid message should not change status or call callbacks
    prev_status = timepix_fly_client._status
    with mock.patch.object(timepix_fly_client, "_run_status_callbacks") as mock_run_callbacks2:
        timepix_fly_client._on_received_ws_message("invalid_status_string")
        # status stays as previous value and callbacks not called
        assert timepix_fly_client._status == prev_status
        mock_run_callbacks2.assert_not_called()


def test_timepix_fly_client_on_status_callbacks(timepix_fly_client):
    """This tests the _run_status_callbacks method of timepix fly client."""
    # Immediate run when current status already in success
    timepix_fly_client._status = TimePixFlyStatus.INIT
    status = StatusBase()
    timepix_fly_client.add_status_callback(
        status=status, success=[TimePixFlyStatus.INIT], error=[TimePixFlyStatus.EXCEPT], run=True
    )
    assert status.done is True and status.success is True

    # Add callback (do not run immediately) and then trigger via _run_status_callbacks
    status2 = StatusBase()
    timepix_fly_client.add_status_callback(
        status=status2,
        success=[TimePixFlyStatus.CONFIG],
        error=[TimePixFlyStatus.EXCEPT],
        run=False,
    )
    # Set status to CONFIG and mark started True to check reset
    timepix_fly_client._status = TimePixFlyStatus.CONFIG
    timepix_fly_client._started = True
    timepix_fly_client._run_status_callbacks()
    assert status2.done is True and status2.success is True
    assert timepix_fly_client._started is False

    # Error path: add with run True when status is EXCEPT
    timepix_fly_client._status = TimePixFlyStatus.EXCEPT
    status3 = StatusBase()
    with mock.patch.object(timepix_fly_client, "last_error") as mock_last_error:
        mock_err = mock.Mock()
        mock_err.message = "boom"
        mock_last_error.return_value = mock_err

        timepix_fly_client.add_status_callback(
            status=status3,
            success=[TimePixFlyStatus.INIT],
            error=[TimePixFlyStatus.EXCEPT],
            run=True,
        )
        assert status3.done is True and status3.success is False


def test_timepix_fly_client_shutdown(timepix_fly_client):
    """This tests the shutdown method of timepix fly client."""
    mock_ws = mock.Mock()
    timepix_fly_client.ws_client = mock_ws

    timepix_fly_client.shutdown()

    mock_ws.close.assert_called_once()
    assert timepix_fly_client.ws_client is None
    assert timepix_fly_client._shutdown_event.is_set()


def test_timepix_fly_client_start(timepix_fly_client):
    """This tests the start method of timepix fly client."""
    # The client._get uses requests.get with f"http://{self.rest_url}/{get_cmd}".
    # We mock requests.get and verify it is called with the expected URL and timeout,
    # and that the _started flag is set.
    module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
    with mock.patch(f"{module_path}.requests.get") as mock_get:
        mock_resp = mock.Mock()
        mock_resp.raise_for_status = mock.Mock()
        mock_resp.text = ""
        mock_get.return_value = mock_resp

        timepix_fly_client.start()

        expected_url = f"http://{timepix_fly_client.rest_url}/?start=true"
        mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout)
        assert timepix_fly_client._started is True


def test_timepix_fly_client_stop_collect(timepix_fly_client):
    """This tests the stop_collect method of timepix fly client."""
    module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
    with mock.patch(f"{module_path}.requests.get") as mock_get:
        mock_resp = mock.Mock()
        mock_resp.raise_for_status = mock.Mock()
        mock_resp.text = ""
        mock_get.return_value = mock_resp

        # ensure started is True and then call stop_collect
        timepix_fly_client._started = True
        timepix_fly_client.stop_collect()

        expected_url = f"http://{timepix_fly_client.rest_url}/?stop_collect=true"
        mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout)
        assert timepix_fly_client._started is False


def test_timepix_fly_client_state(timepix_fly_client):
    """This tests the state method of timepix fly client."""
    module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
    with mock.patch(f"{module_path}.requests.get") as mock_get:
        mock_resp = mock.Mock()
        mock_resp.raise_for_status = mock.Mock()
        # Return a JSON payload compatible with ProgramState
        mock_resp.json.return_value = {"type": "ProgramState", "state": "init"}
        mock_get.return_value = mock_resp

        program_state = timepix_fly_client.state()

        expected_url = f"http://{timepix_fly_client.rest_url}/state"
        mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout)
        # ProgramState.model defines 'state' as a string literal; ensure we parsed it
        assert hasattr(program_state, "state")
        assert program_state.state == "init"


def test_timepix_fly_client_set_pixel_map(timepix_fly_client):
    """This tests the set_pixel_map/_put path by mocking requests.put and checking payload."""
    module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
    with mock.patch(f"{module_path}.requests.put") as mock_put:
        mock_resp = mock.Mock()
        mock_resp.raise_for_status = mock.Mock()
        mock_put.return_value = mock_resp

        # Minimal valid PixelMap dict (type is optional; model supplies default)
        pixel_map = {"chips": [[{"i": 0, "p": 1, "f": 2}]]}

        timepix_fly_client.set_pixel_map(pixel_map)

        expected_url = f"http://{timepix_fly_client.rest_url}/pixel-map"
        # Verify requests.put called with expected url, json and timeout
        mock_put.assert_called_once()
        _, kwargs = mock_put.call_args
        assert kwargs.get("timeout") == timepix_fly_client._timeout
        assert kwargs.get("json") is not None
        assert kwargs.get("json").get("chips") == pixel_map["chips"]