"""Tests for net/fault.py: iptables rule builder, apply_partition_rules, remove_rules_by_comment."""

from __future__ import annotations

from unittest.mock import AsyncMock, call, patch

import pytest

from vmcluster_mcp.net.fault import (
    apply_partition_rules,
    build_partition_rules,
    remove_rules_by_comment,
)


# ---------------------------------------------------------------------------
# build_partition_rules — rule generation
# ---------------------------------------------------------------------------


def test_build_partition_rules_1_vs_1_produces_two_rules():
    """1-node vs 1-node partition: A→B drop + B→A drop = 2 rules."""
    rules = build_partition_rules("mycluster", "h1", "virbr0", ["10.0.0.1"], ["10.0.0.2"])
    assert len(rules) == 2


def test_build_partition_rules_2_vs_2_produces_eight_rules():
    """2-node vs 2-node partition: 4 pairs × 2 symmetric directions = 8 rules."""
    rules = build_partition_rules(
        "mycluster", "h2", "virbr0",
        ["10.0.0.1", "10.0.0.2"],
        ["10.0.0.3", "10.0.0.4"],
    )
    assert len(rules) == 8


def test_build_partition_rules_symmetric():
    """Every A→B rule has a corresponding B→A reverse rule."""
    ips_a = ["10.0.0.1"]
    ips_b = ["10.0.0.2"]
    rules = build_partition_rules("c", "h", "br0", ips_a, ips_b)
    # Extract (src, dst) pairs
    pairs = {(r[r.index("-s") + 1], r[r.index("-d") + 1]) for r in rules}
    assert ("10.0.0.1", "10.0.0.2") in pairs
    assert ("10.0.0.2", "10.0.0.1") in pairs


def test_build_partition_rules_comment_tag_format():
    """Every rule must contain the tagged comment vmcluster:<cluster>:<handle_id>."""
    rules = build_partition_rules("testcluster", "abc123", "virbr0", ["1.1.1.1"], ["2.2.2.2"])
    expected_tag = "vmcluster:testcluster:abc123"
    for rule in rules:
        assert "--comment" in rule
        comment_idx = rule.index("--comment")
        assert rule[comment_idx + 1] == expected_tag


def test_build_partition_rules_targets_bridge_interface():
    """Rules must specify -i and -o with the bridge name for FORWARD chain."""
    rules = build_partition_rules("c", "h", "virbr0", ["1.1.1.1"], ["2.2.2.2"])
    for rule in rules:
        assert "-i" in rule
        assert rule[rule.index("-i") + 1] == "virbr0"
        assert "-o" in rule
        assert rule[rule.index("-o") + 1] == "virbr0"


def test_build_partition_rules_targets_forward_chain():
    """Rules must insert into the FORWARD chain."""
    rules = build_partition_rules("c", "h", "br0", ["1.1.1.1"], ["2.2.2.2"])
    for rule in rules:
        assert "-I" in rule
        assert rule[rule.index("-I") + 1] == "FORWARD"


def test_build_partition_rules_action_is_drop():
    """Rules must use -j DROP."""
    rules = build_partition_rules("c", "h", "br0", ["1.1.1.1"], ["2.2.2.2"])
    for rule in rules:
        assert "-j" in rule
        assert rule[rule.index("-j") + 1] == "DROP"


def test_build_partition_rules_uses_iptables():
    """First element of each rule command must be 'iptables'."""
    rules = build_partition_rules("c", "h", "br0", ["1.1.1.1"], ["2.2.2.2"])
    for rule in rules:
        assert rule[0] == "iptables"


# ---------------------------------------------------------------------------
# apply_partition_rules
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_apply_partition_rules_calls_run_host_cmd_for_each_rule():
    """apply_partition_rules invokes _run_host_cmd once per rule."""
    rules = build_partition_rules("c", "h1", "br0", ["1.1.1.1"], ["2.2.2.2"])
    mock_cmd = AsyncMock(return_value=(0, ""))
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        identifiers = await apply_partition_rules(rules)
    assert mock_cmd.call_count == len(rules)


@pytest.mark.asyncio
async def test_apply_partition_rules_returns_comment_tags():
    """apply_partition_rules returns a list of comment tag strings."""
    rules = build_partition_rules("c", "h1", "br0", ["1.1.1.1"], ["2.2.2.2"])
    mock_cmd = AsyncMock(return_value=(0, ""))
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        identifiers = await apply_partition_rules(rules)
    # Should return comment tags (strings)
    assert isinstance(identifiers, list)
    assert len(identifiers) > 0
    for tag in identifiers:
        assert isinstance(tag, str)
        assert tag.startswith("vmcluster:")


@pytest.mark.asyncio
async def test_apply_partition_rules_raises_on_failure():
    """apply_partition_rules raises RuntimeError if any iptables call fails."""
    rules = build_partition_rules("c", "h1", "br0", ["1.1.1.1"], ["2.2.2.2"])
    mock_cmd = AsyncMock(return_value=(1, "iptables: Permission denied"))
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        with pytest.raises(RuntimeError, match="iptables"):
            await apply_partition_rules(rules)


# ---------------------------------------------------------------------------
# remove_rules_by_comment
# ---------------------------------------------------------------------------

_IPTABLES_LIST_WITH_MATCH = """\
Chain FORWARD (policy ACCEPT)
num   pkts bytes target prot opt in   out   source      destination
1      0     0   DROP   all  --  br0  br0   1.1.1.1     2.2.2.2     /* vmcluster:c:abc */
2      0     0   DROP   all  --  br0  br0   2.2.2.2     1.1.1.1     /* vmcluster:c:abc */
3      0     0   ACCEPT all  --  any  any   0.0.0.0/0   0.0.0.0/0
"""

_IPTABLES_LIST_NO_MATCH = """\
Chain FORWARD (policy ACCEPT)
num   pkts bytes target prot opt in   out   source      destination
1      0     0   ACCEPT all  --  any  any   0.0.0.0/0   0.0.0.0/0
"""


@pytest.mark.asyncio
async def test_remove_rules_by_comment_returns_count_removed():
    """Returns the number of rules successfully removed."""
    call_results = [
        (0, _IPTABLES_LIST_WITH_MATCH),  # list call
        (0, ""),  # delete rule 2 (reverse order)
        (0, ""),  # delete rule 1
    ]
    mock_cmd = AsyncMock(side_effect=call_results)
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        count = await remove_rules_by_comment("vmcluster:c:abc")
    assert count == 2


@pytest.mark.asyncio
async def test_remove_rules_by_comment_deletes_in_reverse_order():
    """Rules are deleted in descending line-number order to avoid index shifting."""
    call_results = [
        (0, _IPTABLES_LIST_WITH_MATCH),
        (0, ""),  # first delete should be for the higher line number
        (0, ""),
    ]
    delete_calls: list[list[str]] = []
    mock_cmd = AsyncMock(side_effect=call_results)
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        await remove_rules_by_comment("vmcluster:c:abc")

    # Second and third calls are deletions
    delete_args = [c.args[0] for c in mock_cmd.call_args_list[1:]]
    line_nums = [int(cmd[-1]) for cmd in delete_args]
    assert line_nums == sorted(line_nums, reverse=True)


@pytest.mark.asyncio
async def test_remove_rules_by_comment_returns_zero_when_no_match():
    """Returns 0 when no matching rules are found — no exception raised."""
    mock_cmd = AsyncMock(return_value=(0, _IPTABLES_LIST_NO_MATCH))
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        count = await remove_rules_by_comment("vmcluster:c:abc")
    assert count == 0


@pytest.mark.asyncio
async def test_remove_rules_by_comment_only_lists_once():
    """Only one iptables list invocation is made before deletions."""
    call_results = [
        (0, _IPTABLES_LIST_WITH_MATCH),
        (0, ""),
        (0, ""),
    ]
    mock_cmd = AsyncMock(side_effect=call_results)
    with patch("vmcluster_mcp.net.fault._run_host_cmd", mock_cmd):
        await remove_rules_by_comment("vmcluster:c:abc")
    # First call is the list call
    list_call_cmd = mock_cmd.call_args_list[0].args[0]
    assert "iptables" in list_call_cmd
    assert "-L" in list_call_cmd
    assert "FORWARD" in list_call_cmd
