"""M16 concurrency and race tests.

Races are the failure class unit tests miss most reliably, because a single-threaded test can
satisfy every assertion while the same code loses an update under load. These tests run real
threads against a shared database and assert on the *outcome* — exactly one winner, no duplicate
authoritative object, no lost update — rather than on timing.

SQLite serialises writers, so a lost update here would be a logic error rather than an isolation
one. The PostgreSQL-specific behaviour (deadlock retry, serialisation failure) is characterised in
the live chaos evidence; what these tests pin down is that the *code* claims a single winner.
"""

from __future__ import annotations

import os
import threading
import uuid
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import pytest
from pydantic import SecretStr
from sqlalchemy import create_engine, event, func, select, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import IntegrityError, OperationalError
from sqlalchemy.orm import Session

from modelforge_api.domain.agent_protocol import (
    AGENT_PROTOCOL_CAPABILITIES,
    AgentMetadata,
    EnrollmentRequest,
    EnrollmentTokenCreate,
)
from modelforge_api.persistence.models import (
    AuditChainHead,
    AuditEvent,
    Base,
    ComputeNode,
    NodeCredential,
    NodeEnrollment,
)
from modelforge_api.services.audit import AuditWriter, audit_chain_violations
from modelforge_api.services.invariants import InvariantStatus, check_invariants
from modelforge_api.services.node_agent import (
    AgentAuthenticationError,
    AgentConflictError,
    NodeAgentService,
)
from modelforge_api.settings import Settings


@pytest.fixture
def engine(tmp_path: Path) -> Engine:
    """A file-backed database so each thread contends on a real connection.

    An in-memory SQLite database behind a StaticPool shares one connection between threads, which
    is an API misuse rather than a concurrency test: the contention being measured would be the
    driver's, not the platform's.
    """

    value = create_engine(f"sqlite+pysqlite:///{tmp_path / 'm16.sqlite'}")

    @event.listens_for(value, "connect")
    def _enable_busy_timeout(connection: Any, _record: Any) -> None:
        cursor = connection.cursor()
        cursor.execute("PRAGMA busy_timeout = 15000")
        cursor.close()

    Base.metadata.create_all(value)
    return value


def settings() -> Settings:
    return Settings(
        _env_file=None,
        operator_api_key=SecretStr("m16-operator-key"),
        node_stale_after_seconds=10,
        node_offline_after_seconds=20,
    )


def metadata() -> AgentMetadata:
    return AgentMetadata(
        agent_version="0.1.0",
        protocol_version=1,
        supported_capabilities=AGENT_PROTOCOL_CAPABILITIES,
        started_at=datetime.now(UTC),
    )


def run_concurrently(
    worker: Callable[[int], Any], count: int, *, max_workers: int | None = None
) -> list[Any]:
    """Start every worker at the same moment so they actually contend."""

    barrier = threading.Barrier(count)

    def wrapped(index: int) -> Any:
        barrier.wait(timeout=30)
        return worker(index)

    with ThreadPoolExecutor(max_workers=max_workers or count) as pool:
        return list(pool.map(wrapped, range(count)))


# --------------------------------------------------------------------- enrolment storm


@pytest.mark.parametrize("attempts", [20, 60])
def test_a_single_use_enrolment_token_survives_a_concurrent_storm(
    engine: Engine, attempts: int
) -> None:
    """M15 closed this race for two threads; M16 proves it holds at storm scale."""

    with Session(engine) as session:
        created = NodeAgentService(session, settings()).create_enrollment(
            EnrollmentTokenCreate(display_name="Storm target")
        )
    token = created.enrollment_token

    def attempt(index: int) -> str:
        with Session(engine) as session:
            service = NodeAgentService(session, settings())
            try:
                service.enroll(
                    EnrollmentRequest(
                        enrollment_token=token,
                        identity_key=f"storm-node-{index}",
                        identity_source="persisted_uuid",
                        hostname=f"storm-{index}",
                        display_name=f"storm-{index}",
                        metadata=metadata(),
                    )
                )
                return "ENROLLED"
            except AgentAuthenticationError:
                return "REFUSED"
            except (IntegrityError, OperationalError):
                # A database-level refusal is still a refusal; what matters is that it is not a
                # second successful identity.
                session.rollback()
                return "REFUSED"

    outcomes = run_concurrently(attempt, attempts)
    assert outcomes.count("ENROLLED") == 1, outcomes
    assert outcomes.count("REFUSED") == attempts - 1

    with Session(engine) as session:
        assert session.scalar(select(func.count()).select_from(ComputeNode)) == 1
        active = session.scalar(
            select(func.count()).select_from(NodeCredential).where(NodeCredential.revoked_at.is_(None))
        )
        assert active == 1
        enrollment = session.scalar(select(NodeEnrollment))
        assert enrollment is not None
        assert enrollment.used_at is not None
        assert enrollment.enrolled_node_id is not None
        report = check_invariants(session)
        assert report.violated == 0, [
            item.key for item in report.results if item.status is InvariantStatus.VIOLATED
        ]


def test_managed_postgresql_claim_and_revocation_linearize_on_the_row_lock() -> None:
    """Managed gate for the PostgreSQL row-lock semantics SQLite cannot reproduce.

    Set MODELFORGE_TEST_POSTGRES_URL to an isolated, disposable PostgreSQL database. The test
    creates and removes one UUID-named schema, pauses enrolment after its conditional claim has
    acquired the row lock, then proves a concurrent revocation loses after the claim commits.
    """

    database_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
    if not database_url:
        pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
    if not database_url.startswith(("postgresql://", "postgresql+psycopg://")):
        pytest.fail("MODELFORGE_TEST_POSTGRES_URL must use PostgreSQL")
    if database_url.startswith("postgresql://"):
        database_url = database_url.replace("postgresql://", "postgresql+psycopg://", 1)

    schema = f"test_enrollment_linearization_{uuid.uuid4().hex}"
    admin_engine = create_engine(database_url, pool_pre_ping=True)
    with admin_engine.begin() as connection:
        connection.execute(text(f'CREATE SCHEMA "{schema}"'))
    engine = create_engine(
        database_url,
        pool_pre_ping=True,
        connect_args={"options": f"-csearch_path={schema}"},
    )
    claim_written = threading.Event()
    allow_claim_commit = threading.Event()
    revoke_update_started = threading.Event()

    @event.listens_for(engine, "before_cursor_execute")
    def _observe_revoke_update(
        _connection: Any,
        _cursor: Any,
        statement: str,
        _parameters: Any,
        _context: Any,
        _executemany: bool,
    ) -> None:
        normalized = statement.lower()
        if normalized.startswith("update node_enrollments set revoked_at"):
            revoke_update_started.set()

    try:
        Base.metadata.create_all(engine)
        with Session(engine) as session:
            created = NodeAgentService(session, settings()).create_enrollment(
                EnrollmentTokenCreate(display_name="PostgreSQL row-lock target")
            )
        enrollment_id = uuid.UUID(created.id)

        def claim() -> str:
            with Session(engine) as session:
                service = NodeAgentService(session, settings())
                original_claim = service._claim_enrollment

                def claim_then_pause(
                    *, enrollment_id: uuid.UUID, token_hash: str, claim_now: datetime
                ) -> bool:
                    won = original_claim(
                        enrollment_id=enrollment_id,
                        token_hash=token_hash,
                        claim_now=claim_now,
                    )
                    if won:
                        claim_written.set()
                        if not allow_claim_commit.wait(timeout=30):
                            raise RuntimeError("timed out while holding enrollment claim")
                    return won

                service._claim_enrollment = claim_then_pause  # type: ignore[method-assign]
                response = service.enroll(
                    EnrollmentRequest(
                        enrollment_token=created.enrollment_token,
                        identity_key="postgres-race-node",
                        identity_source="persisted_uuid",
                        hostname="postgres-race-node",
                        display_name="postgres-race-node",
                        metadata=metadata(),
                    )
                )
                return response.credential_id

        def revoke() -> str:
            with Session(engine) as session:
                try:
                    NodeAgentService(session, settings()).revoke_enrollment(enrollment_id)
                except AgentConflictError:
                    return "CONFLICT"
                return "REVOKED"

        with ThreadPoolExecutor(max_workers=2) as pool:
            claim_future = pool.submit(claim)
            revoke_future = None
            try:
                assert claim_written.wait(timeout=30), "claim never acquired its row lock"
                revoke_future = pool.submit(revoke)
                assert revoke_update_started.wait(timeout=30), "revoke UPDATE never started"
                assert not revoke_future.done(), "revoke did not wait for the claim row lock"
            finally:
                allow_claim_commit.set()
            credential_id = claim_future.result(timeout=30)
            assert revoke_future is not None
            assert revoke_future.result(timeout=30) == "CONFLICT"

        with Session(engine) as session:
            enrollment = session.get(NodeEnrollment, enrollment_id)
            credential = session.get(NodeCredential, uuid.UUID(credential_id))
            assert enrollment is not None and enrollment.used_at is not None
            assert enrollment.revoked_at is None
            assert credential is not None and credential.revoked_at is None
            assert (
                session.scalar(
                    select(func.count())
                    .select_from(NodeCredential)
                    .where(NodeCredential.revoked_at.is_(None))
                )
                == 1
            )
            assert (
                session.scalar(
                    select(func.count())
                    .select_from(AuditEvent)
                    .where(AuditEvent.action == "NODE_ENROLLMENT_TOKEN_REVOKED")
                )
                == 0
            )
    finally:
        allow_claim_commit.set()
        engine.dispose()
        with admin_engine.begin() as connection:
            connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
        admin_engine.dispose()


def test_a_storm_of_distinct_tokens_creates_exactly_one_node_each(engine: Engine) -> None:
    """The guard must reject a reused token without rejecting legitimate parallel enrolments."""

    count = 12
    with Session(engine) as session:
        service = NodeAgentService(session, settings())
        tokens = [
            service.create_enrollment(
                EnrollmentTokenCreate(display_name=f"Node {index}")
            ).enrollment_token
            for index in range(count)
        ]

    def attempt(index: int) -> str:
        with Session(engine) as session:
            try:
                NodeAgentService(session, settings()).enroll(
                    EnrollmentRequest(
                        enrollment_token=tokens[index],
                        identity_key=f"parallel-node-{index}",
                        identity_source="persisted_uuid",
                        hostname=f"parallel-{index}",
                        display_name=f"parallel-{index}",
                        metadata=metadata(),
                    )
                )
                return "ENROLLED"
            except (
                AgentAuthenticationError,
                AgentConflictError,
                IntegrityError,
                OperationalError,
            ):
                session.rollback()
                return "REFUSED"

    outcomes = run_concurrently(attempt, count)
    assert outcomes.count("ENROLLED") == count, outcomes

    with Session(engine) as session:
        assert session.scalar(select(func.count()).select_from(ComputeNode)) == count
        assert (
            session.scalar(
                select(func.count())
                .select_from(NodeCredential)
                .where(NodeCredential.revoked_at.is_(None))
            )
            == count
        )
        assert check_invariants(session).violated == 0


def test_re_enrolment_of_one_identity_under_load_keeps_one_active_credential(
    engine: Engine,
) -> None:
    """Repeated re-enrolment of the same hardware must never leave two usable credentials."""

    identity = "gpu_node-hardware"
    rounds = 8
    with Session(engine) as session:
        service = NodeAgentService(session, settings())
        tokens = [
            service.create_enrollment(
                EnrollmentTokenCreate(display_name="GPU Node")
            ).enrollment_token
            for _ in range(rounds)
        ]

    def attempt(index: int) -> str:
        with Session(engine) as session:
            try:
                NodeAgentService(session, settings()).enroll(
                    EnrollmentRequest(
                        enrollment_token=tokens[index],
                        identity_key=identity,
                        identity_source="persisted_uuid",
                        hostname="gpu_node",
                        display_name="GPU Node",
                        metadata=metadata(),
                    )
                )
                return "ENROLLED"
            except (
                AgentAuthenticationError,
                AgentConflictError,
                IntegrityError,
                OperationalError,
            ):
                session.rollback()
                return "REFUSED"

    outcomes = run_concurrently(attempt, rounds)
    assert "ENROLLED" in outcomes

    with Session(engine) as session:
        nodes = list(session.scalars(select(ComputeNode)))
        assert len(nodes) == 1
        assert nodes[0].key == identity
        active = list(
            session.scalars(select(NodeCredential).where(NodeCredential.revoked_at.is_(None)))
        )
        assert len(active) == 1
        report = check_invariants(session)
        assert report.violated == 0


# --------------------------------------------------------------------- idempotency storm


def test_the_same_idempotency_key_in_parallel_produces_one_logical_operation(
    engine: Engine,
) -> None:
    """A retried request under load must converge on one row, not many."""

    from modelforge_api.persistence.models import ServingJob

    key = uuid.uuid4().hex
    deployment = uuid.uuid4()
    node = uuid.uuid4()
    attempts = 24

    def attempt(_index: int) -> str:
        with Session(engine) as session:
            existing = session.scalar(
                select(ServingJob).where(ServingJob.idempotency_key == key)
            )
            if existing is not None:
                return "REUSED"
            session.add(
                ServingJob(
                    capability_deployment_id=deployment,
                    compute_node_id=node,
                    operation="invoke",
                    status="queued",
                    priority="production",
                    idempotency_key=key,
                )
            )
            try:
                session.commit()
                return "CREATED"
            except (IntegrityError, OperationalError):
                session.rollback()
                return "REUSED"

    outcomes = run_concurrently(attempt, attempts)
    assert outcomes.count("CREATED") >= 1

    with Session(engine) as session:
        rows = session.scalar(
            select(func.count()).select_from(ServingJob).where(ServingJob.idempotency_key == key)
        )
        assert rows == 1, f"{outcomes.count('CREATED')} creators produced {rows} rows"


# --------------------------------------------------------------------- audit-chain serialization


def test_concurrent_audit_appends_form_one_strict_chain(engine: Engine) -> None:
    attempts = 24

    def append(index: int) -> int:
        with Session(engine) as session:
            event = AuditWriter(session, "operator", f"writer-{index}").write(
                "CONCURRENT_APPEND",
                "audit-test",
                str(index),
                {"attempt": index},
            )
            session.commit()
            return event.sequence

    sequences = run_concurrently(append, attempts)
    assert sorted(sequences) == list(range(1, attempts + 1))

    with Session(engine) as session:
        events = list(session.scalars(select(AuditEvent).order_by(AuditEvent.sequence)))
        assert len(events) == attempts
        assert audit_chain_violations(events, session.get(AuditChainHead, 1)) == []
        result = next(
            item
            for item in check_invariants(session).results
            if item.key == "audit_chain_intact"
        )
        assert result.status is InvariantStatus.HOLDS


# --------------------------------------------------------------------- credential races


def test_concurrent_revocation_and_authentication_never_accepts_a_revoked_credential(
    engine: Engine,
) -> None:
    """Whichever order the race resolves in, a revoked credential must never authenticate."""

    with Session(engine) as session:
        service = NodeAgentService(session, settings())
        created = service.create_enrollment(EnrollmentTokenCreate(display_name="Race target"))
        response = service.enroll(
            EnrollmentRequest(
                enrollment_token=created.enrollment_token,
                identity_key="race-node",
                identity_source="persisted_uuid",
                hostname="race",
                display_name="race",
                metadata=metadata(),
            )
        )
        credential = response.node_credential
        node_id = session.scalar(select(ComputeNode.id))

    results: list[str] = []
    lock = threading.Lock()

    def authenticate(_index: int) -> None:
        with Session(engine) as session:
            try:
                NodeAgentService(session, settings()).authenticate(f"Bearer {credential}")
                outcome = "ACCEPTED"
            except AgentAuthenticationError:
                outcome = "REFUSED"
            except (IntegrityError, OperationalError):
                session.rollback()
                outcome = "REFUSED"
        with lock:
            results.append(outcome)

    def revoke(_index: int) -> None:
        with Session(engine) as session:
            try:
                NodeAgentService(session, settings()).revoke_credential(node_id)
            except (IntegrityError, OperationalError):
                session.rollback()

    def worker(index: int) -> None:
        if index == 0:
            revoke(index)
        else:
            authenticate(index)

    run_concurrently(worker, 16)

    with Session(engine) as session:
        stored = session.scalar(select(NodeCredential))
        assert stored is not None
        assert stored.revoked_at is not None
        with pytest.raises(AgentAuthenticationError):
            NodeAgentService(session, settings()).authenticate(f"Bearer {credential}")
        report = check_invariants(session)
        assert report.violated == 0


def test_authentication_after_revocation_is_refused_for_every_attempt(engine: Engine) -> None:
    """No window exists in which a revoked credential is intermittently accepted."""

    with Session(engine) as session:
        service = NodeAgentService(session, settings())
        created = service.create_enrollment(EnrollmentTokenCreate(display_name="Revoked target"))
        response = service.enroll(
            EnrollmentRequest(
                enrollment_token=created.enrollment_token,
                identity_key="revoked-node",
                identity_source="persisted_uuid",
                hostname="revoked",
                display_name="revoked",
                metadata=metadata(),
            )
        )
        credential = response.node_credential
        service.revoke_credential(session.scalar(select(ComputeNode.id)))

    def attempt(_index: int) -> str:
        with Session(engine) as session:
            try:
                NodeAgentService(session, settings()).authenticate(f"Bearer {credential}")
                return "ACCEPTED"
            except AgentAuthenticationError:
                return "REFUSED"

    outcomes = run_concurrently(attempt, 24)
    assert set(outcomes) == {"REFUSED"}, outcomes


# --------------------------------------------------------------------- invariants under load


def test_invariants_stay_readable_while_the_database_is_being_written(engine: Engine) -> None:
    """The safety check must not need a quiet system to answer."""

    stop = threading.Event()

    def writer() -> None:
        index = 0
        while not stop.is_set() and index < 200:
            with Session(engine) as session:
                session.add(ComputeNode(key=f"writer-{index}", hostname=f"writer-{index}"))
                try:
                    session.commit()
                except (IntegrityError, OperationalError):
                    session.rollback()
            index += 1

    thread = threading.Thread(target=writer)
    thread.start()
    try:
        for _ in range(15):
            with Session(engine) as session:
                report = check_invariants(session)
                assert report.checked == 16
                assert report.violated == 0
    finally:
        stop.set()
        thread.join(timeout=30)