"""Tests for SSH worker invocation and local artifact restoration."""

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace
from typing import Any


def test_runtime_uses_workflow_host_snapshot_without_global_registry(
    ssh_executor_module: Any,
    local_ui_events_module: Any,
    remote_hosts_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A queued workflow host should be sufficient to prepare its assigned worker."""
    host = remote_hosts_module.SshHostConfig(
        environment_id="workflow-host",
        display_name="Workflow host",
        ssh_target="workflow-host",
        maximum_workers=2,
    )
    observed: dict[str, Any] = {}

    class FakeController:
        """Capture the rehydrated workflow host."""

        def __init__(self, configured_host: Any) -> None:
            """Record the host passed to the controller."""
            observed["host"] = configured_host

        def probe_capabilities(self) -> None:
            """Record the compatibility reprobe for a legacy host snapshot."""
            observed["probed"] = True
            return None

    class FakeManager:
        """Capture the requested worker slot."""

        def __init__(self, **kwargs: Any) -> None:
            """Record construction arguments."""
            observed["manager"] = kwargs

        def ensure_worker(
            self,
            worker_index: int,
            status_callback: Any = None,
        ) -> Any:
            """Return a deterministic worker specification."""
            observed["worker_index"] = worker_index
            observed["status_callback"] = status_callback
            status_callback("Building SSH runtime environment=workflow-host")
            return SimpleNamespace(worker_index=worker_index)

    monkeypatch.setattr(ssh_executor_module, "SshDockerController", FakeController)
    monkeypatch.setattr(ssh_executor_module, "SshRuntimeManager", FakeManager)
    monkeypatch.setattr(
        local_ui_events_module,
        "_emit_local_remote_startup_status",
        lambda payload, **kwargs: observed.setdefault("status_events", []).append(
            (payload, kwargs)
        ),
    )
    client = ssh_executor_module.SshDockerExecutorClient(
        registry=None,
        repo_root=tmp_path,
        settings=SimpleNamespace(),
    )

    _manager, spec = client._runtime(
        {
            "execution_environment_id": "workflow-host",
            "execution_worker_index": 1,
            "ssh_host_config": host.to_dict(),
        }
    )

    assert observed["host"] == host
    assert observed["probed"] is True
    assert observed["worker_index"] == 1
    assert callable(observed["status_callback"])
    assert observed["status_events"] == [
        (
            {
                "execution_environment_id": "workflow-host",
                "execution_worker_index": 1,
                "ssh_host_config": host.to_dict(),
            },
            {
                "phase": "starting",
                "status_message": "Building SSH runtime environment=workflow-host",
            },
        )
    ]
    assert spec.worker_index == 1


def test_ssh_stream_materializes_remote_output_artifacts(
    ssh_executor_module: Any,
    remote_modal_app_module: Any,
    output_artifacts_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """SSH results should restore output files before deserializing node outputs."""
    output_directory = tmp_path / "output"
    output_directory.mkdir()
    packed_result = output_artifacts_module.pack_remote_execution_result(
        output_artifacts_module.RemoteExecutionResult(
            outputs=b"serialized outputs",
            artifacts=(
                output_artifacts_module.RemoteOutputArtifact(
                    relative_path="images/result.png",
                    payload=b"png bytes",
                ),
            ),
            completed_epoch=1_700_000_000,
        )
    )
    settings = SimpleNamespace(
        comfyui_root=tmp_path,
        app_name="test-app",
    )
    client = ssh_executor_module.SshDockerExecutorClient(
        registry=SimpleNamespace(),
        repo_root=tmp_path,
        settings=settings,
    )
    manager = SimpleNamespace()
    spec = SimpleNamespace()
    monkeypatch.setattr(client, "_runtime", lambda _payload: (manager, spec))
    monkeypatch.setattr(
        client,
        "_invoke_stream",
        lambda *_args: iter(({"kind": "result", "outputs": packed_result},)),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_consume_remote_payload_stream",
        lambda _payload, stream: next(stream)["outputs"],
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_comfy_output_directory",
        lambda _settings: output_directory,
    )

    outputs = client._consume_stream({"component_id": "1"}, b"inputs")

    assert outputs == b"serialized outputs"
    materialized = list(output_directory.rglob("*.png"))
    assert len(materialized) == 1
    assert materialized[0].read_bytes() == b"png bytes"


def test_ssh_transport_failure_restarts_worker_and_retries_once(
    ssh_executor_module: Any,
    remote_modal_app_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A broken relay should recover the exact worker slot with the stable request."""
    stopped: list[int] = []
    manager = SimpleNamespace(
        controller=SimpleNamespace(host=SimpleNamespace(environment_id="host")),
        stop_worker=lambda worker_index: stopped.append(worker_index),
    )
    spec = SimpleNamespace(worker_index=2)
    client = ssh_executor_module.SshDockerExecutorClient(
        registry=SimpleNamespace(),
        repo_root=tmp_path,
        settings=SimpleNamespace(comfyui_root=tmp_path, app_name="app"),
    )
    attempts = 0

    def invoke(*_args: Any) -> Any:
        """Fail the first relay attempt and complete the second."""
        nonlocal attempts
        attempts += 1
        if attempts == 1:
            raise ssh_executor_module.SshRemoteTransportError("connection lost")
        return iter(({"kind": "result", "outputs": b"serialized"},))

    monkeypatch.setattr(client, "_runtime", lambda _payload: (manager, spec))
    monkeypatch.setattr(client, "_invoke_stream", invoke)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_consume_remote_payload_stream",
        lambda _payload, stream: next(stream)["outputs"],
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_materialize_remote_execution_result",
        lambda response, settings: response,
    )

    result = client._consume_stream({"component_id": "1"}, b"inputs")

    assert result == b"serialized"
    assert attempts == 2
    assert stopped == [2]


def test_ssh_dispatch_stages_and_rewrites_hugging_face_model_reference(
    ssh_executor_module: Any,
    remote_modal_app_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """SSH workers should execute immutable generated profiles, not raw model IDs."""
    requested_model = "owner/model"
    generated_profile = "hf-" + "a" * 64
    payload = {
        "component_id": "llm",
        "subgraph_prompt": {
            "llm": {
                "class_type": "ModalLLM",
                "inputs": {"model_profile": requested_model},
            }
        },
    }
    manager = SimpleNamespace(
        controller=SimpleNamespace(host=SimpleNamespace(environment_id="lambda"))
    )
    spec = SimpleNamespace(worker_index=0, container_name="worker")
    client = ssh_executor_module.SshDockerExecutorClient(
        registry=SimpleNamespace(),
        repo_root=tmp_path,
        settings=SimpleNamespace(comfyui_root=tmp_path, app_name="app"),
    )
    stage_calls: list[list[str]] = []

    def stage(*_args: Any) -> list[dict[str, Any]]:
        """Return one generated immutable SSH profile."""
        model_references = _args[-1]
        stage_calls.append(model_references)
        return [
            {
                "requested_reference": requested_model,
                "profile_id": generated_profile,
                "revision": "7" * 40,
            }
        ]

    def invoke(_manager: Any, _spec: Any, rewritten: Any, _inputs: bytes) -> Any:
        """Assert that only the generated profile reaches the worker executor."""
        assert (
            rewritten["subgraph_prompt"]["llm"]["inputs"]["model_profile"]
            == generated_profile
        )
        return iter(({"kind": "result", "outputs": b"serialized"},))

    with ssh_executor_module._STAGED_SSH_PROFILES_LOCK:
        ssh_executor_module._STAGED_SSH_PROFILE_RESULTS.clear()
    monkeypatch.setattr(client, "_runtime", lambda _payload: (manager, spec))
    monkeypatch.setattr(client, "_run_profile_stager", stage)
    monkeypatch.setattr(client, "_invoke_stream", invoke)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_consume_remote_payload_stream",
        lambda _payload, stream: next(stream)["outputs"],
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_materialize_remote_execution_result",
        lambda response, settings: response,
    )

    result = client._consume_stream(payload, b"inputs")

    assert result == b"serialized"
    assert stage_calls == [[requested_model]]


def test_ssh_cancel_targets_owner_during_model_staging(
    ssh_executor_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Cancellation must target one stager without stopping its worker."""
    manager = SimpleNamespace(
        controller=SimpleNamespace(host=SimpleNamespace(environment_id="lambda")),
    )
    spec = SimpleNamespace(worker_index=3)
    process = SimpleNamespace()
    invocation_id = "RIV_staging"
    client = ssh_executor_module.SshDockerExecutorClient(
        registry=SimpleNamespace(),
        repo_root=tmp_path,
        settings=SimpleNamespace(comfyui_root=tmp_path, app_name="app"),
    )
    with ssh_executor_module._ACTIVE_SSH_STAGERS_LOCK:
        ssh_executor_module._ACTIVE_SSH_STAGERS[invocation_id] = (
            manager,
            spec,
            "ssh:RIV_staging:owner",
            process,
        )
    targeted: list[tuple[Any, Any, str]] = []
    terminated: list[Any] = []
    monkeypatch.setattr(
        ssh_executor_module,
        "terminate_staging_transport",
        terminated.append,
    )
    client._cancel_remote_stager = (  # type: ignore[method-assign]
        lambda active_manager, active_spec, owner_id, **_kwargs: (
            targeted.append((active_manager, active_spec, owner_id)) or True
        )
    )
    try:
        cancelled = client.cancel({}, invocation_id)
    finally:
        with ssh_executor_module._ACTIVE_SSH_STAGERS_LOCK:
            ssh_executor_module._ACTIVE_SSH_STAGERS.pop(invocation_id, None)

    assert cancelled is True
    assert terminated == [process]
    assert targeted == [(manager, spec, "ssh:RIV_staging:owner")]