from __future__ import annotations

import uuid
from datetime import UTC, datetime

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from pydantic import SecretStr
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
from sqlalchemy.pool import StaticPool

from modelforge_api.api.routes.agent import (
    get_agent_service,
    get_decommission_service,
)
from modelforge_api.api.routes.agent import (
    router as agent_router,
)
from modelforge_api.domain.agent_protocol import (
    AgentMetadata,
    HeartbeatRequest,
    NodeManagementUpdate,
)
from modelforge_api.domain.node_decommission import NodeDecommissionExecute
from modelforge_api.domain.node_lifecycle import (
    NodeLifecycleChange,
    NodeLifecyclePreviewRequest,
)
from modelforge_api.main import agent_protocol_error, node_decommission_error
from modelforge_api.persistence.models import (
    Accelerator,
    AuditEvent,
    Base,
    ComputeNode,
    GatewayRequest,
    GpuLease,
    NodeCredential,
)
from modelforge_api.services.node_agent import (
    AgentAuthorizationError,
    AgentConflictError,
    NodeAgentService,
    secret_hash,
)
from modelforge_api.services.node_decommission import (
    NodeDecommissionError,
    NodeDecommissionService,
)
from modelforge_api.settings import Settings, get_settings


@pytest.fixture
def session() -> Session:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(engine)
    with Session(engine) as value:
        yield value


def _node(session: Session, *, enabled: bool = True, liveness: str = "offline") -> ComputeNode:
    node = ComputeNode(
        key=f"node-{uuid.uuid4()}",
        hostname="node-lifecycle-test",
        display_name="Lifecycle test node",
        enabled=enabled,
        status="active",
        liveness_state=liveness,
        inventory={"source": "test"},
        observation_source="remote_agent",
        agent_protocol_version=1,
        agent_capabilities=["hardware.inventory"],
        agent_version="1.3.0-rc.1",
        last_heartbeat_at=datetime.now(UTC),
    )
    session.add(node)
    session.commit()
    return node


def _request(preview, *, target_enabled: bool | None = None) -> NodeLifecycleChange:
    return NodeLifecycleChange(
        target_enabled=preview.target_enabled if target_enabled is None else target_enabled,
        expected_generation=preview.node_generation,
        dependency_digest=preview.dependency_digest,
        operator="test-operator",
        reason="Planned node maintenance with workloads drained.",
        confirmation=preview.persisted_identity,
    )


def test_disable_requires_preview_confirmation_and_audits_transition(session: Session) -> None:
    node = _node(session, liveness="online")
    service = NodeDecommissionService(session)

    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
    )
    assert preview.can_apply is True
    assert preview.current_enabled is True
    assert preview.target_enabled is False
    assert preview.blockers == []

    result = service.apply_enabled_change(node.id, _request(preview))

    session.refresh(node)
    assert result.enabled is False
    assert result.liveness == "disabled"
    assert node.enabled is False
    assert node.liveness_state == "disabled"
    event = session.scalar(
        select(AuditEvent).where(
            AuditEvent.action == "NODE_DISABLED",
            AuditEvent.resource_id == str(node.id),
        )
    )
    assert event is not None
    assert event.details["reason"] == "Planned node maintenance with workloads drained."


def test_disable_fails_closed_for_active_gateway_request_and_gpu_lease(session: Session) -> None:
    node = _node(session)
    accelerator = Accelerator(
        compute_node_id=node.id,
        device_index=0,
        device_uuid="GPU-node-lifecycle",
        name="Test GPU",
    )
    session.add(accelerator)
    session.flush()
    gateway = GatewayRequest(
        request_id=uuid.uuid4(),
        capability_key="assistant.general",
        capability_version=1,
        compute_node_id=node.id,
        status="running",
        priority="production",
        input_sha256="a" * 64,
        input_count=1,
    )
    lease = GpuLease(
        accelerator_id=accelerator.id,
        deployment_id=uuid.uuid4(),
        priority="production",
        reserved_vram_mb=1024,
        state="active",
    )
    session.add_all([gateway, lease])
    session.commit()
    service = NodeDecommissionService(session)

    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
    )
    assert preview.can_apply is False
    assert {item.code for item in preview.blockers} >= {
        "active_gateway_request",
        "active_gpu_lease",
    }
    with pytest.raises(NodeDecommissionError) as raised:
        service.apply_enabled_change(node.id, _request(preview))
    assert raised.value.code == "node_lifecycle_blocked"
    session.refresh(node)
    assert node.enabled is True


def test_disable_rejects_preview_when_work_arrives_after_impact_review(session: Session) -> None:
    node = _node(session)
    service = NodeDecommissionService(session)
    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
    )
    request = _request(preview)
    session.add(
        GatewayRequest(
            request_id=uuid.uuid4(),
            capability_key="assistant.general",
            capability_version=1,
            compute_node_id=node.id,
            status="queued",
            priority="production",
            input_sha256="b" * 64,
            input_count=1,
        )
    )
    session.commit()

    with pytest.raises(NodeDecommissionError) as raised:
        service.apply_enabled_change(node.id, request)
    assert raised.value.code == "node_lifecycle_preview_stale"
    session.refresh(node)
    assert node.enabled is True
    assert node.liveness_state == "offline"


def test_node_generation_change_invalidates_lifecycle_preview(session: Session) -> None:
    node = _node(session)
    service = NodeDecommissionService(session)
    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
    )
    request = _request(preview)
    node.display_name = "Renamed after preview"
    session.commit()

    with pytest.raises(NodeDecommissionError) as raised:
        service.apply_enabled_change(node.id, request)
    assert raised.value.code == "node_generation_conflict"
    session.refresh(node)
    assert node.enabled is True


def test_enable_requires_active_remote_credential_and_preserves_unknown_readiness(
    session: Session,
) -> None:
    node = _node(session, enabled=False, liveness="disabled")
    session.add(
        NodeCredential(
            compute_node_id=node.id,
            secret_hash="c" * 64,
            scope="node.publish",
        )
    )
    session.commit()
    service = NodeDecommissionService(session)

    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=True)
    )
    assert preview.can_apply is True
    result = service.apply_enabled_change(node.id, _request(preview))

    session.refresh(node)
    assert result.enabled is True
    assert result.liveness == "offline"
    assert result.agent_evidence_state == "awaiting_fresh_authenticated_heartbeat"
    assert node.enabled is True
    assert node.liveness_state == "offline"
    assert node.agent_protocol_version is None
    assert node.agent_capabilities == []
    assert node.agent_version is None
    assert node.last_heartbeat_at is None
    assert "awaiting fresh authenticated heartbeat" in node.last_connection_error


def test_enable_without_active_credential_and_enable_of_tombstone_are_blocked(
    session: Session,
) -> None:
    node = _node(session, enabled=False, liveness="disabled")
    service = NodeDecommissionService(session)
    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=True)
    )
    assert preview.can_apply is False
    assert "active_node_credential_missing" in {item.code for item in preview.blockers}
    with pytest.raises(NodeDecommissionError) as raised:
        service.apply_enabled_change(node.id, _request(preview))
    assert raised.value.code == "node_lifecycle_blocked"
    with pytest.raises(AgentConflictError, match="impact preview and typed confirmation"):
        NodeAgentService(session, Settings(_env_file=None)).update_node(
            node.id, NodeManagementUpdate(enabled=True)
        )


def test_direct_node_patch_rejects_enabled_but_keeps_metadata_updates(session: Session) -> None:
    node = _node(session)
    service = NodeAgentService(session, Settings(_env_file=None))

    updated = service.update_node(node.id, NodeManagementUpdate(display_name="Renamed safely"))
    assert updated.display_name == "Renamed safely"
    with pytest.raises(AgentConflictError, match="impact preview and typed confirmation"):
        service.update_node(node.id, NodeManagementUpdate(enabled=False))
    session.refresh(node)
    assert node.enabled is True


def test_agent_heartbeat_and_post_body_auth_refuse_node_disabled_after_prebody_auth(
    session: Session,
) -> None:
    node = _node(session)
    credential_id = uuid.uuid4()
    raw_credential = f"mfnode_{credential_id}_race-test-secret"
    session.add(
        NodeCredential(
            id=credential_id,
            compute_node_id=node.id,
            secret_hash=secret_hash(raw_credential),
            scope="node.publish",
        )
    )
    session.commit()

    # Keep the pre-body identity cached like a request whose lifecycle changes while its body is
    # being read. A later admin transaction disables the same row before request dispatch.
    agent_session = Session(session.get_bind(), 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_credential}"
    )
    agent_session.commit()
    assert stale_node.enabled is True

    lifecycle_session = Session(session.get_bind())
    lifecycle = NodeDecommissionService(lifecycle_session)
    preview = lifecycle.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
    )
    lifecycle.apply_enabled_change(node.id, _request(preview))
    lifecycle_session.close()

    with pytest.raises(AgentAuthorizationError, match="compute node is disabled"):
        agent_service.reuse_authentication(evidence, f"Bearer {raw_credential}")

    heartbeat = HeartbeatRequest(
        identity_key=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),
    )
    with pytest.raises(AgentAuthorizationError, match="compute node is disabled"):
        agent_service.heartbeat(stale_node, heartbeat)

    session.refresh(node)
    assert node.enabled is False
    assert node.liveness_state == "disabled"
    agent_session.close()


def test_decommissioned_identity_cannot_be_reenabled(session: Session) -> None:
    node = _node(session, enabled=False, liveness="disabled")
    service = NodeDecommissionService(session)
    retirement = service.preview(node.id)
    service.execute(
        node.id,
        NodeDecommissionExecute(
            expected_generation=retirement.node_generation,
            preview_digest=retirement.dependency_digest,
            idempotency_key=f"decom-{uuid.uuid4()}",
            operator="test-operator",
            reason="This test node has left service permanently.",
            confirmation=retirement.persisted_identity,
        ),
    )

    preview = service.preview_enabled_change(
        node.id, NodeLifecyclePreviewRequest(target_enabled=True)
    )
    assert preview.can_apply is False
    assert "node_decommissioned" in {item.code for item in preview.blockers}
    with pytest.raises(NodeDecommissionError) as raised:
        service.apply_enabled_change(node.id, _request(preview))
    assert raised.value.code == "node_lifecycle_blocked"


def test_admin_api_requires_preview_for_enabled_patch_but_preserves_metadata_patch(
    session: Session,
) -> None:
    node = _node(session)
    settings = Settings(_env_file=None, operator_api_key=SecretStr("admin-key"))
    agent = NodeAgentService(session, settings)
    lifecycle = NodeDecommissionService(session)
    test_app = FastAPI()
    test_app.include_router(agent_router)
    test_app.add_exception_handler(AgentConflictError, agent_protocol_error)
    test_app.add_exception_handler(NodeDecommissionError, node_decommission_error)
    test_app.dependency_overrides[get_agent_service] = lambda: agent
    test_app.dependency_overrides[get_decommission_service] = lambda: lifecycle
    test_app.dependency_overrides[get_settings] = lambda: settings

    try:
        with TestClient(test_app) as client:
            path = f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle/preview"
            assert client.post(path, json={"target_enabled": False}).status_code == 401
            headers = {"X-ModelForge-Admin-Token": "admin-key"}
            preview_response = client.post(path, headers=headers, json={"target_enabled": False})
            assert preview_response.status_code == 200
            preview = preview_response.json()
            assert preview["can_apply"] is True

            bypass = client.patch(
                f"/api/v1/admin/hardware/nodes/{node.id}",
                headers=headers,
                json={"enabled": False},
            )
            assert bypass.status_code == 409
            assert bypass.json()["error"]["code"] == "agent_conflict"
            metadata = client.patch(
                f"/api/v1/admin/hardware/nodes/{node.id}",
                headers=headers,
                json={"display_name": "Updated metadata"},
            )
            assert metadata.status_code == 200
            assert node.display_name == "Updated metadata"

            stale_change = {
                "target_enabled": False,
                "expected_generation": preview["node_generation"],
                "dependency_digest": preview["dependency_digest"],
                "operator": "test-operator",
                "reason": "Planned node maintenance after draining workloads.",
                "confirmation": preview["persisted_identity"],
            }
            stale = client.post(
                f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle",
                headers=headers,
                json=stale_change,
            )
            assert stale.status_code == 409
            assert stale.json()["error"]["code"] == "node_generation_conflict"

            refreshed_preview = client.post(
                path, headers=headers, json={"target_enabled": False}
            ).json()
            change = {**stale_change, "expected_generation": refreshed_preview["node_generation"], "dependency_digest": refreshed_preview["dependency_digest"]}
            applied = client.post(
                f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle",
                headers=headers,
                json=change,
            )
            assert applied.status_code == 200
            assert applied.json()["enabled"] is False
            assert applied.json()["liveness"] == "disabled"
    finally:
        test_app.dependency_overrides.clear()