"""Python nodes run in a worker process, and stay there when things go wrong."""

import sys
import threading
import time
from collections.abc import Iterator

import pytest

from app.flow.workers import NodeCancelled, NodeTimeout, PythonWorkerPool


@pytest.fixture
def pool() -> Iterator[PythonWorkerPool]:
    # One worker: a respawn is then provably the same slot coming back.
    worker_pool = PythonWorkerPool(python=sys.executable, size=1)
    worker_pool.start()
    yield worker_pool
    worker_pool.stop()


def run(pool: PythonWorkerPool, code: str, node: str = "demo", **kwargs):
    return pool.run(
        "demo", node, code, kwargs, {"factor": 2}, f"demo.{node}", timeout=5
    )


def test_a_node_returns_its_value_and_what_it_printed(pool, capsys):
    result = run(
        pool,
        "def process(value, params):\n"
        "    print('seen', value)\n"
        "    return {'out': value * params['factor']}\n",
        value=21,
    )
    assert result == {"out": 42}
    # The proxy writes them to stdout, which is where the engine's tee is.
    assert "seen 21" in capsys.readouterr().out


def test_a_failure_keeps_its_class_and_points_at_the_node(pool):
    with pytest.raises(Exception) as caught:
        run(pool, "def process(params):\n    raise ValueError('bad input')\n")

    # The engine renders a node error as "<class>: <message>", so both have to
    # survive the trip.
    assert type(caught.value).__name__ == "ValueError"
    assert str(caught.value) == "bad input"
    assert "<node demo." in caught.value.remote_traceback
    assert "ValueError: bad input" in caught.value.remote_traceback


def test_a_node_that_kills_its_worker_is_an_ordinary_error(pool):
    with pytest.raises(Exception, match="worker died"):
        run(pool, "import os\n\n\ndef process(params):\n    os._exit(1)\n")

    assert run(pool, "def process(params):\n    return {'out': 1}\n") == {"out": 1}


def test_a_node_that_runs_too_long_is_killed_and_the_pool_recovers(pool):
    started = time.monotonic()
    with pytest.raises(NodeTimeout):
        pool.run(
            "demo",
            "slow",
            "import time\n\n\ndef process(params):\n    time.sleep(30)\n",
            {},
            {},
            "demo.slow",
            timeout=1,
        )
    assert time.monotonic() - started < 10

    # The killed worker's slot is refilled on the next call.
    assert run(pool, "def process(params):\n    return {'out': 2}\n") == {"out": 2}


def test_a_running_node_can_be_cancelled(pool):
    def stop_it() -> None:
        for _ in range(100):
            if pool.cancel("demo.slow"):
                return
            time.sleep(0.05)

    stopper = threading.Thread(target=stop_it)
    stopper.start()
    try:
        with pytest.raises(NodeCancelled):
            pool.run(
                "demo",
                "slow",
                "import time\n\n\ndef process(params):\n    time.sleep(30)\n",
                {},
                {},
                "demo.slow",
                timeout=30,
            )
    finally:
        stopper.join()


def test_a_result_that_is_not_json_is_refused(pool):
    with pytest.raises(Exception, match="cannot be sent back as JSON"):
        run(pool, "def process(params):\n    return {'out': {1, 2}}\n")


def test_compiling_reports_where_the_source_is_wrong(pool):
    assert pool.compile("demo", "broken", "def process(params)\n    return {}\n")
    assert pool.compile("demo", "fine", "def process(params):\n    return {}\n") is None


def test_a_node_imports_the_standard_library_not_the_engines_own_modules(pool):
    # The worker script lives in app/flow, which holds queue.py, secrets.py and
    # more; the interpreter would put that directory first on sys.path.
    result = run(
        pool,
        "import queue\nimport secrets\n\n\n"
        "def process(params):\n"
        "    return {'out': [queue.Queue().qsize(), len(secrets.token_hex(4))]}\n",
    )
    assert result == {"out": [0, 8]}


def test_the_engines_secrets_are_not_in_a_workers_environment(pool, monkeypatch):
    monkeypatch.setenv("SECRET_KEY", "not-for-nodes")
    monkeypatch.setenv("POSTGRES_PASSWORD", "not-for-nodes")
    monkeypatch.setenv("FLUKSIO_HARMLESS", "fine")
    # A fresh process, so it is built from the environment set just now.
    pool.respawn_all()

    result = run(
        pool,
        "import os\n\n\n"
        "def process(params):\n"
        "    return {'out': [k for k in ('SECRET_KEY', 'POSTGRES_PASSWORD',\n"
        "                                'FLUKSIO_HARMLESS') if k in os.environ]}\n",
    )
    assert result == {"out": ["FLUKSIO_HARMLESS"]}


def test_a_pool_can_stop_while_a_node_is_running(pool):
    # One slot, taken by a node that will not finish on its own, and a second
    # call queued behind it. The engine's node threads are not daemons, so a
    # wait here is a shutdown that never completes.
    outcomes: list[str] = []

    def call(node: str) -> None:
        try:
            pool.run(
                "demo",
                node,
                "import time\n\n\ndef process(params):\n    time.sleep(60)\n",
                {},
                {},
                f"demo.{node}",
                timeout=60,
            )
            outcomes.append("returned")
        except Exception as exc:
            outcomes.append(type(exc).__name__)

    busy = threading.Thread(target=call, args=("busy",))
    busy.start()
    # Let the first one take the slot, so the second is blocked acquiring it.
    time.sleep(1)
    waiting = threading.Thread(target=call, args=("waiting",))
    waiting.start()
    time.sleep(0.2)

    pool.stop()
    for thread in (busy, waiting):
        thread.join(timeout=10)
        assert not thread.is_alive()
    assert len(outcomes) == 2

    with pytest.raises(Exception, match="shutting down"):
        run(pool, "def process(params):\n    return {'out': 1}\n")


# -----------------------------------------------------------------------------
# Reporting from inside a node that has not returned yet
# -----------------------------------------------------------------------------


def test_a_node_reports_metrics_while_it_is_still_running(pool):
    seen = []
    result = pool.run(
        "demo",
        "train",
        "import fluksio\n"
        "def process(params):\n"
        "    for step in range(3):\n"
        "        fluksio.log_metric('loss', 1.0 / (step + 1), step)\n"
        "    fluksio.progress(0.5, 'halfway')\n"
        "    return {'out': 1}\n",
        {},
        {},
        "demo.train",
        timeout=5,
        run_id="r1",
        on_event=seen.append,
    )

    assert result == {"out": 1}
    metrics = [event for event in seen if event["event"] == "metric"]
    assert [(m["name"], m["step"]) for m in metrics] == [
        ("loss", 0),
        ("loss", 1),
        ("loss", 2),
    ]
    assert metrics[0]["value"] == 1.0
    # Every event says which call it belongs to, so a sweep can tell them apart.
    assert {m["call_id"] for m in metrics} == {"r1:demo.train"}
    assert [event["event"] for event in seen if event["event"] == "progress"] == [
        "progress"
    ]


def test_events_hold_off_the_timeout_but_silence_does_not(pool):
    # The deadline measures silence: a node reporting every 0.05s stays alive
    # well past a 0.3s timeout, which is what a two-hour training needs.
    result = pool.run(
        "demo",
        "slow",
        "import time, fluksio\n"
        "def process(params):\n"
        "    for step in range(12):\n"
        "        time.sleep(0.05)\n"
        "        fluksio.log_metric('beat', step, step)\n"
        "    return {'done': True}\n",
        {},
        {},
        "demo.slow",
        timeout=0.3,
        run_id="r2",
        on_event=lambda _event: None,
    )
    assert result == {"done": True}

    with pytest.raises(NodeTimeout):
        pool.run(
            "demo",
            "quiet",
            "import time\ndef process(params):\n    time.sleep(2)\n    return {}\n",
            {},
            {},
            "demo.quiet",
            timeout=0.3,
        )


def test_cancelling_one_run_leaves_the_same_node_in_another_alone(pool):
    # Keyed by (run, node): cancelling a config of a sweep must not kill the
    # rest of it. With one slot the second run is not executing, so the check
    # is that the pool refuses to find it rather than killing the wrong worker.
    started = threading.Event()

    def hold():
        try:
            pool.run(
                "demo",
                "hold",
                "import time\ndef process(params):\n    time.sleep(5)\n    return {}\n",
                {},
                {},
                "demo.hold",
                timeout=10,
                run_id="run-a",
            )
        except Exception:
            pass
        finally:
            started.set()

    thread = threading.Thread(target=hold, daemon=True)
    thread.start()
    time.sleep(0.5)

    assert pool.cancel("demo.hold", run_id="run-b") is False
    assert pool.cancel("demo.hold", run_id="run-a") is True
    started.wait(timeout=5)
    thread.join(timeout=5)