"""Tests for brynq.swarm_state_machine — rigid state machine for swarm lifecycle."""

import json
import time
import pytest
from pathlib import Path


@pytest.fixture(autouse=True)
def _isolate_swarm(tmp_path, monkeypatch):
    """Redirect swarm session storage to a temp directory."""
    sessions_dir = tmp_path / "swarm_sessions"
    sessions_dir.mkdir()
    history_path = sessions_dir / "history.jsonl"
    monkeypatch.setattr("brynq.swarm_state_machine.SESSIONS_DIR", sessions_dir)
    monkeypatch.setattr("brynq.swarm_state_machine.HISTORY_PATH", history_path)


# ── create_swarm_session ───────────────────────────────────────────────────


class TestCreateSwarmSession:
    def test_returns_session_id(self):
        from brynq.swarm_state_machine import create_swarm_session
        sid = create_swarm_session("Test Swarm", "Build a feature")
        assert isinstance(sid, str)
        assert len(sid) == 12

    def test_initial_state_is_planning(self):
        from brynq.swarm_state_machine import create_swarm_session, get_state
        sid = create_swarm_session("Test", "Task")
        assert get_state(sid) == "PLANNING"

    def test_creates_json_file(self):
        from brynq.swarm_state_machine import create_swarm_session, SESSIONS_DIR
        sid = create_swarm_session("Test", "Task")
        assert (SESSIONS_DIR / f"{sid}.json").exists()

    def test_session_has_correct_fields(self):
        from brynq.swarm_state_machine import create_swarm_session, _read_session
        sid = create_swarm_session("My Swarm", "Do things", required_skills=["python"], max_agents=3)
        s = _read_session(sid)
        assert s["name"] == "My Swarm"
        assert s["task_description"] == "Do things"
        assert s["required_skills"] == ["python"]
        assert s["max_agents"] == 3
        assert s["agents"] == {}
        assert s["progress"] == {}

    def test_default_max_agents(self):
        from brynq.swarm_state_machine import create_swarm_session, _read_session
        sid = create_swarm_session("Test", "Task")
        s = _read_session(sid)
        assert s["max_agents"] == 5

    def test_logs_creation_event(self):
        from brynq.swarm_state_machine import create_swarm_session, get_session_history
        sid = create_swarm_session("Test", "Task")
        history = get_session_history(sid)
        assert len(history) == 1
        assert history[0]["event_type"] == "created"

    def test_empty_name_raises(self):
        from brynq.swarm_state_machine import create_swarm_session
        with pytest.raises(ValueError, match="Name"):
            create_swarm_session("", "Task")

    def test_empty_description_raises(self):
        from brynq.swarm_state_machine import create_swarm_session
        with pytest.raises(ValueError, match="Task description"):
            create_swarm_session("Name", "  ")

    def test_zero_max_agents_raises(self):
        from brynq.swarm_state_machine import create_swarm_session
        with pytest.raises(ValueError, match="max_agents"):
            create_swarm_session("Name", "Task", max_agents=0)


# ── get_state ──────────────────────────────────────────────────────────────


class TestGetState:
    def test_returns_current_state(self):
        from brynq.swarm_state_machine import create_swarm_session, get_state
        sid = create_swarm_session("Test", "Task")
        assert get_state(sid) == "PLANNING"

    def test_nonexistent_raises(self):
        from brynq.swarm_state_machine import get_state
        with pytest.raises(FileNotFoundError):
            get_state("nonexistent")


# ── transition ─────────────────────────────────────────────────────────────


class TestTransition:
    def test_valid_planning_to_recruiting(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        assert get_state(sid) == "RECRUITING"

    def test_valid_recruiting_to_executing(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        assert get_state(sid) == "EXECUTING"

    def test_valid_executing_to_reviewing(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        transition(sid, "REVIEWING")
        assert get_state(sid) == "REVIEWING"

    def test_valid_reviewing_to_completed(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        transition(sid, "REVIEWING")
        transition(sid, "COMPLETED")
        assert get_state(sid) == "COMPLETED"

    def test_valid_executing_to_blocked(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        transition(sid, "BLOCKED")
        assert get_state(sid) == "BLOCKED"

    def test_valid_blocked_to_executing(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        transition(sid, "BLOCKED")
        transition(sid, "EXECUTING")
        assert get_state(sid) == "EXECUTING"

    def test_valid_reviewing_to_executing(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        transition(sid, "REVIEWING")
        transition(sid, "EXECUTING")
        assert get_state(sid) == "EXECUTING"

    def test_invalid_skip_states(self):
        from brynq.swarm_state_machine import create_swarm_session, transition
        sid = create_swarm_session("Test", "Task")
        with pytest.raises(ValueError, match="Invalid transition"):
            transition(sid, "EXECUTING")

    def test_invalid_planning_to_completed(self):
        from brynq.swarm_state_machine import create_swarm_session, transition
        sid = create_swarm_session("Test", "Task")
        with pytest.raises(ValueError, match="Invalid transition"):
            transition(sid, "COMPLETED")

    def test_cannot_transition_from_terminal(self):
        from brynq.swarm_state_machine import create_swarm_session, transition
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        transition(sid, "REVIEWING")
        transition(sid, "COMPLETED")
        with pytest.raises(ValueError, match="terminal state"):
            transition(sid, "EXECUTING")

    def test_any_state_can_fail(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_state
        sid = create_swarm_session("Test", "Task")
        transition(sid, "FAILED")
        assert get_state(sid) == "FAILED"

    def test_invalid_state_name(self):
        from brynq.swarm_state_machine import create_swarm_session, transition
        sid = create_swarm_session("Test", "Task")
        with pytest.raises(ValueError, match="Invalid state"):
            transition(sid, "BOGUS")

    def test_nonexistent_session(self):
        from brynq.swarm_state_machine import transition
        with pytest.raises(FileNotFoundError):
            transition("nonexistent", "RECRUITING")

    def test_logs_transition_event(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_session_history
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING", reason="Ready to recruit")
        history = get_session_history(sid)
        trans_events = [e for e in history if e["event_type"] == "transition"]
        assert len(trans_events) == 1
        assert trans_events[0]["from_state"] == "PLANNING"
        assert trans_events[0]["to_state"] == "RECRUITING"
        assert trans_events[0]["reason"] == "Ready to recruit"


# ── assign_agent / remove_agent / get_agents ───────────────────────────────


class TestAgentManagement:
    def test_assign_agent(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, get_agents
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "agent-1", role="lead")
        agents = get_agents(sid)
        assert len(agents) == 1
        assert agents[0]["agent_id"] == "agent-1"
        assert agents[0]["role"] == "lead"

    def test_assign_default_role(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, get_agents
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "agent-1")
        agents = get_agents(sid)
        assert agents[0]["role"] == "worker"

    def test_assign_multiple_agents(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, get_agents
        sid = create_swarm_session("Test", "Task", max_agents=3)
        assign_agent(sid, "a1", role="lead")
        assign_agent(sid, "a2", role="worker")
        assign_agent(sid, "a3", role="reviewer")
        agents = get_agents(sid)
        assert len(agents) == 3

    def test_assign_exceeds_max_agents(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent
        sid = create_swarm_session("Test", "Task", max_agents=1)
        assign_agent(sid, "a1")
        with pytest.raises(ValueError, match="1/1"):
            assign_agent(sid, "a2")

    def test_assign_duplicate_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        with pytest.raises(ValueError, match="already assigned"):
            assign_agent(sid, "a1")

    def test_assign_invalid_role_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent
        sid = create_swarm_session("Test", "Task")
        with pytest.raises(ValueError, match="Invalid role"):
            assign_agent(sid, "a1", role="boss")

    def test_assign_to_terminal_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, assign_agent
        sid = create_swarm_session("Test", "Task")
        transition(sid, "FAILED")
        with pytest.raises(ValueError, match="FAILED"):
            assign_agent(sid, "a1")

    def test_remove_agent(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, remove_agent, get_agents
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        assign_agent(sid, "a2")
        remove_agent(sid, "a1")
        agents = get_agents(sid)
        assert len(agents) == 1
        assert agents[0]["agent_id"] == "a2"

    def test_remove_nonexistent_agent_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, remove_agent
        sid = create_swarm_session("Test", "Task")
        with pytest.raises(ValueError, match="not in this session"):
            remove_agent(sid, "ghost")

    def test_remove_clears_progress(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, transition, assign_agent,
            report_progress, remove_agent, _read_session,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        transition(sid, "RECRUITING")
        transition(sid, "EXECUTING")
        report_progress(sid, "a1", 50.0)
        remove_agent(sid, "a1")
        s = _read_session(sid)
        assert "a1" not in s["progress"]


# ── report_progress ────────────────────────────────────────────────────────


class TestReportProgress:
    def test_reports_progress(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, report_progress, _read_session,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        report_progress(sid, "a1", 42.5, "Working on it")
        s = _read_session(sid)
        assert s["progress"]["a1"]["pct"] == 42.5
        assert s["progress"]["a1"]["message"] == "Working on it"

    def test_progress_out_of_range_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, report_progress
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        with pytest.raises(ValueError, match="progress_pct"):
            report_progress(sid, "a1", 101.0)

    def test_progress_negative_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, report_progress
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        with pytest.raises(ValueError, match="progress_pct"):
            report_progress(sid, "a1", -1.0)

    def test_progress_nonmember_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, report_progress
        sid = create_swarm_session("Test", "Task")
        with pytest.raises(ValueError, match="not in this session"):
            report_progress(sid, "ghost", 50.0)

    def test_progress_in_terminal_state_raises(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, transition, report_progress,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        transition(sid, "FAILED")
        with pytest.raises(ValueError, match="FAILED"):
            report_progress(sid, "a1", 50.0)

    def test_get_agents_includes_progress(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, report_progress, get_agents,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        report_progress(sid, "a1", 75.0, "Almost done")
        agents = get_agents(sid)
        assert agents[0]["progress"]["pct"] == 75.0


# ── check_blocked ──────────────────────────────────────────────────────────


class TestCheckBlocked:
    def test_not_blocked_when_no_agents(self):
        from brynq.swarm_state_machine import create_swarm_session, check_blocked
        sid = create_swarm_session("Test", "Task")
        result = check_blocked(sid)
        assert result["blocked"] is False

    def test_detects_stuck_agent(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, report_progress,
            check_blocked, _read_session, _write_session,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        report_progress(sid, "a1", 10.0)
        # Manually set progress epoch to the past
        s = _read_session(sid)
        s["progress"]["a1"]["epoch"] = time.time() - 3600  # 1 hour ago
        _write_session(s)
        result = check_blocked(sid)
        assert result["blocked"] is True
        assert "a1" in result["stuck_agents"]

    def test_agent_without_progress_is_stuck(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, check_blocked
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        result = check_blocked(sid)
        assert result["blocked"] is True
        assert "a1" in result["stuck_agents"]

    def test_not_blocked_in_terminal_state(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, transition, check_blocked,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        transition(sid, "FAILED")
        result = check_blocked(sid)
        assert result["blocked"] is False

    def test_nonexistent_raises(self):
        from brynq.swarm_state_machine import check_blocked
        with pytest.raises(FileNotFoundError):
            check_blocked("nonexistent")


# ── auto_reassign ──────────────────────────────────────────────────────────


class TestAutoReassign:
    def test_reassigns_stuck_agent(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, report_progress,
            auto_reassign, _read_session, _write_session,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "stuck-agent")
        assign_agent(sid, "good-agent")
        report_progress(sid, "good-agent", 50.0)
        # Make stuck-agent stuck
        s = _read_session(sid)
        s["progress"]["stuck-agent"] = {"pct": 10, "message": "", "updated_at": "", "epoch": time.time() - 3600}
        _write_session(s)
        result = auto_reassign(sid)
        assert len(result["reassigned"]) == 1
        assert result["reassigned"][0]["from_agent"] == "stuck-agent"
        assert result["reassigned"][0]["to_agent"] == "good-agent"

    def test_no_reassign_when_not_blocked(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, report_progress, auto_reassign,
        )
        sid = create_swarm_session("Test", "Task")
        assign_agent(sid, "a1")
        report_progress(sid, "a1", 50.0)
        result = auto_reassign(sid)
        assert result["reassigned"] == []

    def test_no_available_agent_for_reassign(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, assign_agent, auto_reassign, _read_session, _write_session,
        )
        sid = create_swarm_session("Test", "Task", max_agents=1)
        assign_agent(sid, "only-agent")
        # Only-agent is stuck (no progress reported)
        result = auto_reassign(sid)
        assert len(result["reassigned"]) == 1
        assert result["reassigned"][0]["to_agent"] is None

    def test_terminal_state_returns_empty(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, transition, auto_reassign,
        )
        sid = create_swarm_session("Test", "Task")
        transition(sid, "FAILED")
        result = auto_reassign(sid)
        assert result["reassigned"] == []


# ── get_session_history ────────────────────────────────────────────────────


class TestGetSessionHistory:
    def test_history_includes_all_events(self):
        from brynq.swarm_state_machine import (
            create_swarm_session, transition, assign_agent, get_session_history,
        )
        sid = create_swarm_session("Test", "Task")
        transition(sid, "RECRUITING")
        assign_agent(sid, "a1")
        history = get_session_history(sid)
        event_types = [e["event_type"] for e in history]
        assert "created" in event_types
        assert "transition" in event_types
        assert "agent_assigned" in event_types

    def test_history_isolated_per_session(self):
        from brynq.swarm_state_machine import create_swarm_session, get_session_history
        sid1 = create_swarm_session("S1", "T1")
        sid2 = create_swarm_session("S2", "T2")
        h1 = get_session_history(sid1)
        h2 = get_session_history(sid2)
        assert all(e["session_id"] == sid1 for e in h1)
        assert all(e["session_id"] == sid2 for e in h2)


# ── list_sessions ──────────────────────────────────────────────────────────


class TestListSessions:
    def test_empty(self):
        from brynq.swarm_state_machine import list_sessions
        assert list_sessions() == []

    def test_lists_created_sessions(self):
        from brynq.swarm_state_machine import create_swarm_session, list_sessions
        create_swarm_session("S1", "T1")
        create_swarm_session("S2", "T2")
        sessions = list_sessions()
        assert len(sessions) == 2

    def test_filters_by_state(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, list_sessions
        sid1 = create_swarm_session("S1", "T1")
        sid2 = create_swarm_session("S2", "T2")
        transition(sid1, "RECRUITING")
        planning = list_sessions(state="PLANNING")
        assert len(planning) == 1
        assert planning[0]["state"] == "PLANNING"

    def test_respects_limit(self):
        from brynq.swarm_state_machine import create_swarm_session, list_sessions
        for i in range(10):
            create_swarm_session(f"S{i}", f"T{i}")
        sessions = list_sessions(limit=3)
        assert len(sessions) == 3

    def test_invalid_state_raises(self):
        from brynq.swarm_state_machine import list_sessions
        with pytest.raises(ValueError, match="Invalid state"):
            list_sessions(state="BOGUS")


# ── dissolve ───────────────────────────────────────────────────────────────


class TestDissolve:
    def test_dissolve_sets_failed(self):
        from brynq.swarm_state_machine import create_swarm_session, dissolve, get_state
        sid = create_swarm_session("Test", "Task")
        dissolve(sid, reason="No longer needed")
        assert get_state(sid) == "FAILED"

    def test_dissolve_records_reason(self):
        from brynq.swarm_state_machine import create_swarm_session, dissolve, _read_session
        sid = create_swarm_session("Test", "Task")
        dissolve(sid, reason="Budget cut")
        s = _read_session(sid)
        assert s["dissolve_reason"] == "Budget cut"
        assert s["dissolved_at"] is not None

    def test_dissolve_terminal_raises(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, dissolve
        sid = create_swarm_session("Test", "Task")
        transition(sid, "FAILED")
        with pytest.raises(ValueError, match="terminal state"):
            dissolve(sid)

    def test_dissolve_logs_event(self):
        from brynq.swarm_state_machine import create_swarm_session, dissolve, get_session_history
        sid = create_swarm_session("Test", "Task")
        dissolve(sid)
        history = get_session_history(sid)
        dissolved = [e for e in history if e["event_type"] == "dissolved"]
        assert len(dissolved) == 1

    def test_nonexistent_raises(self):
        from brynq.swarm_state_machine import dissolve
        with pytest.raises(FileNotFoundError):
            dissolve("nonexistent")


# ── get_swarm_stats ────────────────────────────────────────────────────────


class TestGetSwarmStats:
    def test_empty_stats(self):
        from brynq.swarm_state_machine import get_swarm_stats
        stats = get_swarm_stats()
        assert stats["total_sessions"] == 0
        assert stats["active_sessions"] == 0

    def test_counts_by_state(self):
        from brynq.swarm_state_machine import create_swarm_session, transition, get_swarm_stats
        s1 = create_swarm_session("S1", "T1")
        s2 = create_swarm_session("S2", "T2")
        s3 = create_swarm_session("S3", "T3")
        transition(s1, "FAILED")
        stats = get_swarm_stats()
        assert stats["total_sessions"] == 3
        assert stats["active_sessions"] == 2
        assert stats["by_state"]["PLANNING"] == 2
        assert stats["by_state"]["FAILED"] == 1

    def test_avg_agents(self):
        from brynq.swarm_state_machine import create_swarm_session, assign_agent, get_swarm_stats
        s1 = create_swarm_session("S1", "T1")
        s2 = create_swarm_session("S2", "T2")
        assign_agent(s1, "a1")
        assign_agent(s1, "a2")
        assign_agent(s2, "a3")
        stats = get_swarm_stats()
        assert stats["avg_agents_per_swarm"] == 1.5
