"""U11-16: shards 1..N-1 of scripts/generate.py must not run beside a live heavy job."""
import importlib.util
import os

import pytest

from align.exit_codes import PreflightRefused

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
INODE = 12345


def _load():
    spec = importlib.util.spec_from_file_location(
        "generate_under_test", os.path.join(REPO, "scripts", "generate.py")
    )
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


def _locks(pid):
    return f"1: FLOCK  ADVISORY  WRITE {pid} 103:06:{INODE} 0 EOF\n2: POSIX  ADVISORY  READ 99 fd:00:777 0 EOF\n"


def _cmd(table):
    return lambda pid: table[pid]


def test_non_zero_shard_refuses_when_a_trainer_holds_the_lock():
    g = _load()
    with pytest.raises(PreflightRefused, match="heavy"):
        g.verify_shard_lock(
            rank=1, lock_inode=INODE, locks_text=_locks(500),
            cmdline_of=_cmd({500: "python scripts/grpo.py configs/x.toml"}),
        )


def test_non_zero_shard_allows_its_own_shard0_sibling_and_a_free_lock():
    g = _load()
    g.verify_shard_lock(
        rank=2, lock_inode=INODE, locks_text=_locks(500),
        cmdline_of=_cmd({500: "python scripts/generate.py --shard-index 0 --num-shards 4"}),
    )
    g.verify_shard_lock(rank=3, lock_inode=INODE, locks_text="", cmdline_of=_cmd({}))


def test_rank0_and_no_heavy_lock_are_untouched():
    g = _load()
    # rank 0 takes the real lock in apply_resource_guards; nothing to verify here
    g.verify_shard_lock(rank=0, lock_inode=INODE, locks_text=_locks(500), cmdline_of=_cmd({500: "x"}))


@pytest.mark.parametrize("idx,n", [(4, 4), (-1, 4), (0, 0), (5, 2)])
def test_shard_bounds_validated(idx, n):
    g = _load()
    with pytest.raises(ValueError, match="shard"):
        g.validate_shards(idx, n)


def test_valid_shard_bounds():
    g = _load()
    g.validate_shards(0, 1)
    g.validate_shards(3, 4)