from __future__ import annotations

from datetime import datetime, timedelta, timezone

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import Session

from govoplan_core.core.runtime_coordination import (
    DistributedLease,
    RuntimeIdentity,
    RuntimeNode,
    RuntimeNodeState,
    StaleFence,
    acquire_lease,
    assert_lease_fence,
    heartbeat_runtime_node,
    list_runtime_nodes,
    register_runtime_node,
    request_runtime_node_drain,
)
from govoplan_core.db.base import Base


def _session() -> tuple[Session, object]:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(
        engine,
        tables=[RuntimeNode.__table__, DistributedLease.__table__],
    )
    return Session(engine, expire_on_commit=False), engine


def _identity(*, node_id: str, incarnation: str) -> RuntimeIdentity:
    return RuntimeIdentity(
        installation_id="installation-1",
        node_id=node_id,
        incarnation=incarnation,
        role="worker",
        software_version="0.1.14",
        composition_hash="a" * 64,
        queues=("default",),
    )


def test_node_incarnation_guards_heartbeat_and_drain_state() -> None:
    session, engine = _session()
    now = datetime(2026, 8, 1, 10, 0, tzinfo=timezone.utc)
    first = _identity(node_id="worker-1", incarnation="first")
    replacement = _identity(node_id="worker-1", incarnation="replacement")
    try:
        register_runtime_node(session, first, now=now)
        heartbeat_runtime_node(session, first, now=now + timedelta(seconds=5))
        register_runtime_node(session, replacement, now=now + timedelta(seconds=10))

        with pytest.raises(RuntimeError, match="stale node incarnation"):
            heartbeat_runtime_node(session, first, now=now + timedelta(seconds=15))

        node = request_runtime_node_drain(
            session,
            installation_id="installation-1",
            node_id="worker-1",
            reason="rolling update",
            now=now + timedelta(seconds=20),
        )
        assert node.state == RuntimeNodeState.DRAINING.value
        rows = list_runtime_nodes(
            session,
            installation_id="installation-1",
            stale_after_seconds=60,
            now=now + timedelta(seconds=25),
        )
        assert rows[0]["incarnation"] == "replacement"
        assert rows[0]["state"] == "draining"
        assert rows[0]["stale"] is False
    finally:
        session.close()
        engine.dispose()


def test_expired_lease_reassignment_increments_fence() -> None:
    session, engine = _session()
    now = datetime(2026, 8, 1, 10, 0, tzinfo=timezone.utc)
    try:
        first = acquire_lease(
            session,
            installation_id="installation-1",
            resource_key="migration:release",
            holder_node_id="node-a",
            holder_incarnation="a1",
            ttl_seconds=30,
            now=now,
        )
        assert first is not None and first.fencing_token == 1
        blocked = acquire_lease(
            session,
            installation_id="installation-1",
            resource_key="migration:release",
            holder_node_id="node-b",
            holder_incarnation="b1",
            ttl_seconds=30,
            now=now + timedelta(seconds=10),
        )
        assert blocked is None

        replacement = acquire_lease(
            session,
            installation_id="installation-1",
            resource_key="migration:release",
            holder_node_id="node-b",
            holder_incarnation="b1",
            ttl_seconds=30,
            now=now + timedelta(seconds=31),
        )
        assert replacement is not None
        assert replacement.fencing_token == 2
        with pytest.raises(StaleFence):
            assert_lease_fence(
                session,
                first,
                now=now + timedelta(seconds=32),
            )
        assert_lease_fence(
            session,
            replacement,
            now=now + timedelta(seconds=32),
        )
    finally:
        session.close()
        engine.dispose()
