"""Who gets the machine: accounting, exclusivity, and what the worker is told.

The failure this exists for is not subtle — five concurrent nodes each sizing a
thread pool to every core starved the engine's own event loop, and three GPU
processes each preallocating most of the card deadlocked at zero utilisation.
Both come down to arithmetic nobody was doing, so the arithmetic is what is
checked here.
"""

import os
import threading
import time

import pytest

from fluksio.flow.resources import (
    THREAD_VARS,
    Allocation,
    ResourceAccountant,
    derive_env,
    fair_share_env,
)
from fluksio.flow.schemas import NodeDef, Resources


def test_what_is_free_is_what_was_handed_out():
    accountant = ResourceAccountant(cpus=4, gpus=0)
    with accountant.claim(Resources(cpus=3)):
        assert accountant.snapshot()["cpus"] == {"total": 4, "free": 1}
    assert accountant.snapshot()["cpus"] == {"total": 4, "free": 4}


def test_a_node_waits_until_there_is_room():
    """Blocking is the mechanism — the same backpressure a worker slot applies."""
    accountant = ResourceAccountant(cpus=2)
    running = threading.Event()
    started = threading.Event()

    def second() -> None:
        with accountant.claim(Resources(cpus=2), node="study.b"):
            started.set()

    with accountant.claim(Resources(cpus=2), node="study.a"):
        thread = threading.Thread(target=second)
        thread.start()
        # Long enough to have run if nothing was holding it back.
        assert not started.wait(0.2)
        assert accountant.snapshot()["waiting"][0]["node"] == "study.b"
        running.set()

    thread.join(timeout=5)
    assert started.is_set()


def test_a_gpu_is_held_by_one_node_at_a_time():
    """The deadlock was three processes each preallocating most of one card."""
    accountant = ResourceAccountant(cpus=8, gpus=2)
    with accountant.claim(Resources(cpus=1, gpus=1)) as first:
        with accountant.claim(Resources(cpus=1, gpus=1)) as second:
            assert set(first.gpus) & set(second.gpus) == set()
            assert accountant.snapshot()["gpus"] == {"total": 2, "free": 0}
    assert accountant.snapshot()["gpus"] == {"total": 2, "free": 2}


def test_asking_for_more_than_there_is_gets_what_there_is():
    """A flow written on a big box still has to run on a laptop."""
    accountant = ResourceAccountant(cpus=2, gpus=0)
    with accountant.claim(Resources(cpus=64, gpus=4)) as allocation:
        assert allocation.cpus == 2
        assert allocation.gpus == ()


def test_everything_comes_back_when_a_node_fails():
    accountant = ResourceAccountant(cpus=4, gpus=1)
    with pytest.raises(ValueError):
        with accountant.claim(Resources(cpus=4, gpus=1)):
            raise ValueError("the node raised")

    assert accountant.snapshot()["cpus"]["free"] == 4
    assert accountant.snapshot()["gpus"]["free"] == 1


def test_many_nodes_at_once_all_finish():
    """The accountant must not deadlock under contention; it is on every call."""
    accountant = ResourceAccountant(cpus=4)
    done = []

    def work() -> None:
        with accountant.claim(Resources(cpus=2), node="study.n"):
            time.sleep(0.01)
            done.append(1)

    threads = [threading.Thread(target=work) for _ in range(8)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=10)

    assert len(done) == 8
    assert accountant.snapshot()["cpus"]["free"] == 4


# -----------------------------------------------------------------------------
# What the worker is started with
# -----------------------------------------------------------------------------


def test_the_share_becomes_the_thread_limit():
    env = derive_env(Resources(cpus=3), Allocation(cpus=3))

    assert all(env[var] == "3" for var in THREAD_VARS)
    assert "CUDA_VISIBLE_DEVICES" not in env


def test_a_gpu_node_is_told_which_card_is_its():
    env = derive_env(Resources(cpus=1, gpus=2), Allocation(cpus=1, gpus=(1, 3)))

    assert env["CUDA_VISIBLE_DEVICES"] == "1,3"


def test_what_the_node_asked_for_wins():
    """A declaration outranks what the allocation would imply — it is deliberate."""
    wanted = Resources(
        cpus=4, env={"OMP_NUM_THREADS": "1", "XLA_FLAGS": "--xla_cpu_x=false"}
    )
    env = derive_env(wanted, Allocation(cpus=4))

    assert env["OMP_NUM_THREADS"] == "1"
    assert env["XLA_FLAGS"] == "--xla_cpu_x=false"
    assert env["MKL_NUM_THREADS"] == "4"


def test_the_shared_pool_divides_what_it_has():
    env = fair_share_env(cpus=8, workers=4)

    assert all(env[var] == "2" for var in THREAD_VARS)
    # Never zero, however many workers there are.
    assert fair_share_env(cpus=2, workers=8)["OMP_NUM_THREADS"] == "1"


def test_an_operator_who_set_one_keeps_it(monkeypatch):
    """An explicit value in the engine's environment is an answer, not a default."""
    monkeypatch.setitem(os.environ, "OMP_NUM_THREADS", "2")
    env = fair_share_env(cpus=16, workers=2)

    assert "OMP_NUM_THREADS" not in env
    assert env["MKL_NUM_THREADS"] == "8"


# -----------------------------------------------------------------------------
# Declaration
# -----------------------------------------------------------------------------


def test_declaring_nothing_stays_exactly_as_it_was():
    """Every flow that exists today parses unchanged and is not accounted for."""
    assert NodeDef(id="poll").resources is None


def test_a_misspelled_resource_is_refused():
    with pytest.raises(ValueError):
        Resources(cpu=4)


# -----------------------------------------------------------------------------
# The whole path, through a real worker
# -----------------------------------------------------------------------------


def test_a_declared_node_is_held_to_its_share(tmp_path):
    """Four nodes wanting two cores each, on a machine with four.

    What is checked is the pair: never more than the machine has in flight at
    once, and every one of them told what it was given — the two halves that
    together are the oversubscription this is for. Also that it does not
    deadlock, since the claim is taken before a worker slot and both block.
    """
    import sys

    from fluksio.flow.controller import FlowController, RunContext
    from fluksio.flow.messages import DType, MessageSpec
    from fluksio.flow.schemas import FlowDef
    from fluksio.flow.state import MemoryState
    from fluksio.flow.store import FlowStore
    from fluksio.flow.workers import PythonWorkerPool

    store = FlowStore(tmp_path / "flows")
    store.write_flow(
        FlowDef(
            name="study",
            mode="batch",
            nodes=[
                NodeDef(
                    id="fit",
                    provides=[MessageSpec(name="threads", dtype=DType.INT)],
                    resources=Resources(cpus=2, env={"XLA_FLAGS": "--x=false"}),
                )
            ],
        )
    )
    store.write_node_source(
        "study",
        "fit",
        "import os, time\n\n\ndef process():\n"
        "    time.sleep(0.2)\n"
        "    return {'threads': int(os.environ['OMP_NUM_THREADS'])}\n",
    )

    pool = PythonWorkerPool(python=sys.executable, size=4)
    pool.start()
    accountant = ResourceAccountant(cpus=4)
    controller = FlowController(store, workers=pool, resources=accountant)
    answers: list[object] = []
    in_flight: list[int] = []

    def once(index: int) -> None:
        pipeline = controller.build_run_pipeline(
            store.read_flow("study"),
            state=MemoryState(),
            run=RunContext(run_id=f"r{index}"),
        )
        pipeline.run({})
        answers.append(pipeline.state.get("study.threads"))

    def watch(until: threading.Event) -> None:
        while not until.is_set():
            in_flight.append(4 - int(accountant.snapshot()["cpus"]["free"]))
            time.sleep(0.01)

    finished = threading.Event()
    watcher = threading.Thread(target=watch, args=(finished,))
    watcher.start()
    threads = [threading.Thread(target=once, args=(i,)) for i in range(4)]
    try:
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join(timeout=60)
            assert not thread.is_alive(), "a claim and a worker slot deadlocked"
    finally:
        finished.set()
        watcher.join(timeout=5)
        pool.stop()

    assert answers == [2, 2, 2, 2]
    assert max(in_flight) <= 4
    assert accountant.snapshot()["cpus"]["free"] == 4
    # One environment, so one extra pool however many nodes derived it.
    assert len(pool._children) == 1