from __future__ import annotations

import asyncio
from types import SimpleNamespace

from govoplan_core.core.runtime_coordination import (
    RuntimeCoordinationError,
    RuntimeIdentity,
)
from govoplan_core.server.runtime_agent import RuntimeNodeAgent


class _Session:
    def __enter__(self):
        return self

    def __exit__(self, *_args) -> None:
        return None

    def commit(self) -> None:
        return None


class _Database:
    def SessionLocal(self) -> _Session:
        return _Session()


class _Consumer:
    def __init__(self) -> None:
        self.cancelled: list[str] = []
        self.added: list[str] = []

    def cancel_task_queue(self, queue: str) -> None:
        self.cancelled.append(queue)

    def add_task_queue(self, queue: str) -> None:
        self.added.append(queue)


def test_api_runtime_agent_fails_readiness_on_heartbeat_error() -> None:
    agent = RuntimeNodeAgent(
        settings=SimpleNamespace(
            installation_id="installation-1",
            runtime_role="api",
            runtime_node_id="api-1",
            runtime_heartbeat_seconds=15,
            celery_queues="",
        ),
        software_version="0.1.14",
        module_ids=("core",),
    )
    agent.coordination_healthy = True

    def fail() -> None:
        raise RuntimeCoordinationError("database unavailable")

    agent._heartbeat = fail
    assert asyncio.run(agent._heartbeat_once()) is False
    assert agent.coordination_healthy is False

    agent._heartbeat = lambda: None
    assert asyncio.run(agent._heartbeat_once()) is True
    assert agent.coordination_healthy is True


def test_worker_disables_consumers_without_reclaiming_stale_identity(
    monkeypatch,
) -> None:
    from govoplan_core import celery_app

    identity = RuntimeIdentity(
        installation_id="installation-1",
        node_id="worker-1",
        incarnation="stale-incarnation",
        role="worker",
        software_version="0.1.14",
        composition_hash="a" * 64,
        queues=("default", "mail"),
    )
    consumer = _Consumer()
    original_state = (
        celery_app._worker_identity,
        celery_app._worker_consumer,
        celery_app._worker_draining,
    )
    celery_app._worker_identity = identity
    celery_app._worker_consumer = consumer
    celery_app._worker_draining = False
    monkeypatch.setattr(celery_app, "get_database", lambda: _Database())

    def reject_heartbeat(*_args, **_kwargs):
        raise RuntimeCoordinationError("stale node incarnation")

    monkeypatch.setattr(
        celery_app,
        "heartbeat_runtime_node",
        reject_heartbeat,
    )
    monkeypatch.setattr(
        celery_app,
        "register_runtime_node",
        lambda *_args, **_kwargs: (_ for _ in ()).throw(
            AssertionError("heartbeat failures must not reclaim an identity")
        ),
    )
    try:
        celery_app._heartbeat_worker_runtime()
        assert consumer.cancelled == ["default", "mail"]
        assert celery_app._worker_draining is True

        monkeypatch.setattr(
            celery_app,
            "heartbeat_runtime_node",
            lambda *_args, **_kwargs: SimpleNamespace(state="active"),
        )
        celery_app._heartbeat_worker_runtime()
        assert consumer.added == ["default", "mail"]
        assert celery_app._worker_draining is False
    finally:
        (
            celery_app._worker_identity,
            celery_app._worker_consumer,
            celery_app._worker_draining,
        ) = original_state