from __future__ import annotations

import importlib
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import TimeoutError as FutureTimeoutError
from datetime import UTC, datetime

import pytest
from sqlalchemy import event, select
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session

pytest_plugins = ("tests.test_managed_identities_postgres",)

from modelforge_api.domain.agent_protocol import AgentMetadata, HeartbeatRequest
from modelforge_api.domain.node_lifecycle import NodeLifecycleChange, NodeLifecyclePreviewRequest
from modelforge_api.persistence.models import ComputeNode, NodeCredential
from modelforge_api.services.node_agent import (
    AgentAuthorizationError,
    NodeAgentService,
    NodeAuthenticationEvidence,
    secret_hash,
)
from modelforge_api.services.node_decommission import (
    NodeDecommissionError,
    NodeDecommissionService,
)
from modelforge_api.settings import Settings
from tests.test_managed_identities_postgres import IdentityDatabase


def test_agent_request_waits_for_disable_and_rechecks_locked_node(
    identity_database: IdentityDatabase,
) -> None:
    """PostgreSQL proof that a pre-body auth result cannot outlive a lifecycle disable."""
    importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head")
    engine = identity_database.runtime
    node_id = uuid.uuid4()
    credential_id = uuid.uuid4()
    raw = f"mfnode_{credential_id}_lifecycle-race-secret"
    with Session(engine) as session:
        node = ComputeNode(
            id=node_id,
            key=f"postgres-race-{node_id}",
            hostname="postgres-race-node",
            display_name="PostgreSQL lifecycle race",
            enabled=True,
            status="active",
            liveness_state="offline",
            inventory={"source": "test"},
            observation_source="remote_agent",
        )
        session.add(node)
        session.add(
            NodeCredential(
                id=credential_id,
                compute_node_id=node_id,
                secret_hash=secret_hash(raw),
                scope="node.publish",
            )
        )
        session.commit()

        preview = NodeDecommissionService(session).preview_enabled_change(
            node_id, NodeLifecyclePreviewRequest(target_enabled=False)
        )
        change = NodeLifecycleChange(
            target_enabled=False,
            expected_generation=preview.node_generation,
            dependency_digest=preview.dependency_digest,
            operator="postgres-race-test",
            reason="Prove that disable serializes with agent authentication.",
            confirmation=preview.persisted_identity,
        )

    request_session = Session(engine, expire_on_commit=False)
    agent_service = NodeAgentService(request_session, Settings(_env_file=None))
    (_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
        f"Bearer {raw}"
    )
    request_session.commit()
    assert stale_node.enabled is True

    lifecycle_has_locks = threading.Event()
    allow_lifecycle_to_commit = threading.Event()
    agent_lock_attempted = threading.Event()

    with Session(engine) as lifecycle_session:
        lifecycle_service = NodeDecommissionService(lifecycle_session)
        lock_dependencies = lifecycle_service._lock_dependency_tables

        def hold_lifecycle_locks() -> None:
            lock_dependencies()
            lifecycle_has_locks.set()
            if not allow_lifecycle_to_commit.wait(timeout=20):
                raise RuntimeError("timed out holding lifecycle locks in race test")

        lifecycle_service._lock_dependency_tables = hold_lifecycle_locks  # type: ignore[method-assign]

        @event.listens_for(engine, "before_cursor_execute")
        def observe_agent_lock_attempt(
            _connection: object,
            _cursor: object,
            statement: str,
            _parameters: object,
            _context: object,
            _executemany: bool,
        ) -> None:
            normalized = statement.lower()
            if (
                threading.current_thread().name.startswith("agent-request")
                and "compute_nodes" in normalized
                and "for update" in normalized
            ):
                agent_lock_attempted.set()

        try:
            with ThreadPoolExecutor(max_workers=1, thread_name_prefix="lifecycle") as lifecycle_pool:
                lifecycle_future = lifecycle_pool.submit(
                    lifecycle_service.apply_enabled_change, node_id, change
                )
                assert lifecycle_has_locks.wait(timeout=10), "lifecycle never acquired its locks"
                with ThreadPoolExecutor(max_workers=1, thread_name_prefix="agent-request") as agent_pool:
                    agent_future = agent_pool.submit(
                        _publish_authenticated_heartbeat,
                        agent_service,
                        evidence,
                        raw,
                        stale_node.key,
                    )
                    assert agent_lock_attempted.wait(timeout=10), "agent never requested its node lock"
                    with pytest.raises(FutureTimeoutError):
                        agent_future.result(timeout=0.25)

                    allow_lifecycle_to_commit.set()
                    lifecycle_future.result(timeout=30)
                    assert agent_future.result(timeout=30) == "rejected"
        finally:
            allow_lifecycle_to_commit.set()
            event.remove(engine, "before_cursor_execute", observe_agent_lock_attempt)
            request_session.close()

    with Session(engine) as session:
        persisted = session.scalar(select(ComputeNode).where(ComputeNode.id == node_id))
        assert persisted is not None
        assert persisted.enabled is False
        assert persisted.liveness_state == "disabled"
        assert persisted.generation == change.expected_generation + 1


def test_heartbeat_wins_node_lock_and_invalidates_lifecycle_preview(
    identity_database: IdentityDatabase,
) -> None:
    """A heartbeat serialized before disable advances the CAS generation."""
    importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head")
    engine = identity_database.runtime
    node_id = uuid.uuid4()
    credential_id = uuid.uuid4()
    raw = f"mfnode_{credential_id}_heartbeat-first-race-secret"
    with Session(engine) as session:
        node = ComputeNode(
            id=node_id,
            key=f"postgres-heartbeat-race-{node_id}",
            hostname="postgres-heartbeat-race-node",
            display_name="PostgreSQL heartbeat race",
            enabled=True,
            status="active",
            liveness_state="offline",
            inventory={"source": "test"},
            observation_source="remote_agent",
        )
        session.add(node)
        session.add(
            NodeCredential(
                id=credential_id,
                compute_node_id=node_id,
                secret_hash=secret_hash(raw),
                scope="node.publish",
            )
        )
        session.commit()
        lifecycle = NodeDecommissionService(session)
        preview = lifecycle.preview_enabled_change(
            node_id, NodeLifecyclePreviewRequest(target_enabled=False)
        )
        change = NodeLifecycleChange(
            target_enabled=False,
            expected_generation=preview.node_generation,
            dependency_digest=preview.dependency_digest,
            operator="postgres-race-test",
            reason="Prove that fresh heartbeat evidence invalidates this preview.",
            confirmation=preview.persisted_identity,
        )

    agent_session = Session(engine, expire_on_commit=False)
    agent_service = NodeAgentService(agent_session, Settings(_env_file=None))
    (_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
        f"Bearer {raw}"
    )
    agent_session.commit()
    heartbeat_has_node_lock = threading.Event()
    allow_heartbeat_to_commit = threading.Event()
    lifecycle_lock_attempted = threading.Event()
    original_commit = agent_session.commit

    def hold_heartbeat_commit() -> None:
        heartbeat_has_node_lock.set()
        if not allow_heartbeat_to_commit.wait(timeout=20):
            raise RuntimeError("timed out holding the heartbeat node lock")
        original_commit()

    agent_session.commit = hold_heartbeat_commit  # type: ignore[method-assign]

    @event.listens_for(engine, "before_cursor_execute")
    def observe_lifecycle_lock_attempt(
        _connection: object,
        _cursor: object,
        statement: str,
        _parameters: object,
        _context: object,
        _executemany: bool,
    ) -> None:
        normalized = statement.lower()
        if (
            threading.current_thread().name.startswith("lifecycle-after-heartbeat")
            and "compute_nodes" in normalized
            and "for update" in normalized
        ):
            lifecycle_lock_attempted.set()

    request = HeartbeatRequest(
        identity_key=stale_node.key,
        metadata=AgentMetadata(
            agent_version="1.3.0-rc.1",
            protocol_version=1,
            supported_capabilities=["hardware.inventory"],
            started_at=datetime.now(UTC),
        ),
        observed_at=datetime.now(UTC),
    )

    try:
        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="heartbeat") as agent_pool:
            heartbeat_future = agent_pool.submit(
                _publish_and_hold_authenticated_heartbeat,
                agent_service,
                evidence,
                raw,
                request,
            )
            assert heartbeat_has_node_lock.wait(timeout=10), "heartbeat never acquired node lock"
            with ThreadPoolExecutor(
                max_workers=1, thread_name_prefix="lifecycle-after-heartbeat"
            ) as lifecycle_pool:
                lifecycle_future = lifecycle_pool.submit(
                    _apply_stale_lifecycle_change, engine, node_id, change
                )
                assert lifecycle_lock_attempted.wait(timeout=10), "lifecycle never requested node lock"
                with pytest.raises(FutureTimeoutError):
                    lifecycle_future.result(timeout=0.25)
                allow_heartbeat_to_commit.set()
                heartbeat_future.result(timeout=30)
                assert lifecycle_future.result(timeout=30) == "node_generation_conflict"
    finally:
        allow_heartbeat_to_commit.set()
        event.remove(engine, "before_cursor_execute", observe_lifecycle_lock_attempt)
        agent_session.close()

    with Session(engine) as session:
        persisted = session.scalar(select(ComputeNode).where(ComputeNode.id == node_id))
        assert persisted is not None
        assert persisted.enabled is True
        assert persisted.liveness_state == "online"
        assert persisted.last_heartbeat_at is not None
        assert persisted.generation == change.expected_generation + 1


def _publish_and_hold_authenticated_heartbeat(
    service: NodeAgentService,
    evidence: NodeAuthenticationEvidence,
    raw_credential: str,
    request: HeartbeatRequest,
) -> None:
    _credential, node = service.reuse_authentication(evidence, f"Bearer {raw_credential}")
    service.heartbeat(node, request)


def _apply_stale_lifecycle_change(
    engine: Engine, node_id: uuid.UUID, change: NodeLifecycleChange
) -> str:
    with Session(engine) as session:
        try:
            NodeDecommissionService(session).apply_enabled_change(node_id, change)
        except NodeDecommissionError as exc:
            if exc.code != "node_generation_conflict":
                raise
            return exc.code
    return "applied"


def _publish_authenticated_heartbeat(
    service: NodeAgentService,
    evidence: NodeAuthenticationEvidence,
    raw_credential: str,
    identity_key: str,
) -> str:
    try:
        _credential, node = service.reuse_authentication(
            evidence, f"Bearer {raw_credential}"
        )
        service.heartbeat(
            node,
            HeartbeatRequest(
                identity_key=identity_key,
                metadata=AgentMetadata(
                    agent_version="1.3.0-rc.1",
                    protocol_version=1,
                    supported_capabilities=["hardware.inventory"],
                    started_at=datetime.now(UTC),
                ),
                observed_at=datetime.now(UTC),
            ),
        )
    except AgentAuthorizationError:
        return "rejected"
    return "published"