from dataclasses import asdict

from aare.common.automation_models import (
    AutomationProgress,
    StepState,
    StepStatus,
    WorkflowStateKind,
)
from aare.daq.config import BeamlineConfig


class _FakeRedis:
    def __init__(self):
        self._store: dict[str, str | int] = {}

    def get(self, key: str):
        return self._store.get(key)

    def set(self, key: str, value):
        self._store[key] = value

    def delete(self, key: str):
        self._store.pop(key, None)

    def incr(self, key: str) -> int:
        value = int(self._store.get(key, 0)) + 1
        self._store[key] = value
        return value


def _make_config_with_fake_redis() -> BeamlineConfig:
    cfg = BeamlineConfig.__new__(BeamlineConfig)
    cfg._BeamlineConfig__bl = "testbeamline"
    cfg._BeamlineConfig__client = _FakeRedis()
    return cfg


def test_automation_progress_state_round_trip_dataclass():
    cfg = _make_config_with_fake_redis()

    progress = AutomationProgress(
        current_step="Center",
        steps=[
            StepState(
                step=WorkflowStateKind.MOUNT,
                status=StepStatus.SUCCESS,
                message="Mount complete",
            ),
            StepState(
                step=WorkflowStateKind.LOOP_CENTRE,
                status=StepStatus.RUNNING,
                message="Centering sample",
            ),
            StepState(
                step=WorkflowStateKind.RASTER,
                status=StepStatus.PENDING,
                message="",
            ),
            StepState(
                step=WorkflowStateKind.DATA_COLLECTION,
                status=StepStatus.PENDING,
                message="",
            ),
            StepState(
                step=WorkflowStateKind.FINAL,
                status=StepStatus.PENDING,
                message="",
            ),
        ],
        finished=False,
        success=None,
    )

    written_state = cfg.set_automation_progress_state(progress)

    assert written_state["seq"] == 1
    assert written_state["progress"] == asdict(progress)

    read_state = cfg.get_automation_progress_state()

    assert read_state["seq"] == 1
    assert read_state["progress"] == asdict(progress)


def test_automation_progress_state_seq_increments():
    cfg = _make_config_with_fake_redis()

    first = AutomationProgress(
        current_step="Mount",
        steps=[
            StepState(
                step=WorkflowStateKind.MOUNT,
                status=StepStatus.RUNNING,
                message="Mounting sample",
            )
        ],
        finished=False,
        success=None,
    )
    second = AutomationProgress(
        current_step="Paused/Finished",
        steps=[
            StepState(
                step=WorkflowStateKind.FINAL,
                status=StepStatus.SUCCESS,
                message="Automation complete",
            )
        ],
        finished=True,
        success=True,
    )

    first_state = cfg.set_automation_progress_state(first)
    second_state = cfg.set_automation_progress_state(second)

    assert first_state["seq"] == 1
    assert second_state["seq"] == 2

    final_state = cfg.get_automation_progress_state()
    assert final_state["seq"] == 2
    assert final_state["progress"] == asdict(second)