#!/usr/bin/env python3
"""Issue #7 — a new agent must be packed into any same-cwd workspace that still
has room in its 2xK grid; a fresh workspace is created only when all of them
are full.

The shim's `new-session` used to take the first pane whose cwd matched, which
was always the (full) first workspace, judged it overflow and opened yet
another workspace: six agents came out as 4 + 1 + 1 instead of 4 + 2.

The mock herdr reports real per-workspace geometry when the state carries
``layout_geometry: "2x2"``, so lib_py.layout sees a full workspace as
overflow and a half-empty sibling as splittable.
"""
from __future__ import annotations

import json
import subprocess
from pathlib import Path

CWD_AGENTS = {
    # workspace w1: full 2x2 grid
    "a1": ("w1", "w1:p1"), "a2": ("w1", "w1:p2"), "a3": ("w1", "w1:p3"), "a4": ("w1", "w1:p4"),
}


def _seed(state_file: Path, sandbox: Path, layout: dict) -> None:
    """layout: workspace id -> number of agents already in it (same cwd)."""
    state = json.loads(state_file.read_text())
    state["layout_geometry"] = "2x2"
    agents = {}
    workspaces = []
    for ws_id, count in layout.items():
        workspaces.append({"workspace_id": ws_id, "label": "default", "cwd": str(sandbox)})
        for i in range(1, count + 1):
            agents[f"{ws_id}-agent-{i}"] = {
                "agent": "claude",
                "agent_status": "idle",
                "status": "running",
                "workspace_id": ws_id,
                "pane_id": f"{ws_id}:p{i}",
                "cwd": str(sandbox),
                "buffer": "Ready",
            }
    state["workspaces"] = workspaces
    state["agents"] = agents
    state_file.write_text(json.dumps(state, indent=2))


def _new_session(sandbox: Path, name: str) -> subprocess.CompletedProcess:
    lib = sandbox / ".agents" / "skills" / "lib.sh"
    script = f'''
    source "{lib}"
    _init_herdr_isolation
    herdr new-session -d -s "{name}" -c "{sandbox}" "claude --dangerously-skip-permissions"
    '''
    return subprocess.run(["bash", "-c", script], capture_output=True, text=True, cwd=str(sandbox), timeout=180)


def _calls(state_file: Path, *prefix: str) -> list:
    calls = json.loads(state_file.read_text()).get("calls", [])
    return [c for c in calls if c[:len(prefix)] == list(prefix)]


def test_agent_packs_into_sibling_workspace_with_room(mam_sandbox, mock_herdr, mock_agents):
    """w1 holds 4 agents (full), w2 holds 1: the sixth agent must split w2."""
    state_file = Path(mock_herdr)
    _seed(state_file, mam_sandbox, {"w1": 4, "w2": 1})

    res = _new_session(mam_sandbox, "sixth-creator-claude")
    assert res.returncode == 0, res.stderr + res.stdout

    splits = _calls(state_file, "pane", "split")
    assert splits, "the sixth agent must be placed by splitting a pane"
    assert all("HERDR_WORKSPACE_ID=w2" in c and "--pane" in c and c[c.index("--pane") + 1].startswith("w2:")
               for c in splits), splits
    assert not _calls(state_file, "workspace", "create"), "no third workspace while w2 has room"


def test_new_workspace_only_when_every_candidate_is_full(mam_sandbox, mock_herdr, mock_agents):
    """w1 and w2 both hold 4 agents: only then is a third workspace created."""
    state_file = Path(mock_herdr)
    _seed(state_file, mam_sandbox, {"w1": 4, "w2": 4})

    res = _new_session(mam_sandbox, "ninth-creator-claude")
    assert res.returncode == 0, res.stderr + res.stdout

    assert not _calls(state_file, "pane", "split"), "a full grid must not be split"
    assert len(_calls(state_file, "workspace", "create")) == 1


def test_first_workspace_with_room_still_wins(mam_sandbox, mock_herdr, mock_agents):
    """Regression guard: with a single half-empty workspace the agent lands
    there, exactly as before."""
    state_file = Path(mock_herdr)
    _seed(state_file, mam_sandbox, {"w1": 2})

    res = _new_session(mam_sandbox, "third-creator-claude")
    assert res.returncode == 0, res.stderr + res.stdout

    splits = _calls(state_file, "pane", "split")
    assert len(splits) == 1 and "HERDR_WORKSPACE_ID=w1" in splits[0], splits
    assert not _calls(state_file, "workspace", "create")