from nebula_cert_manager.firewall import resolve_firewall_rules
from nebula_cert_manager.models import (
    Direction,
    FirewallConfig,
    Protocol,
    SecurityGroup,
)


def _make_config() -> FirewallConfig:
    return FirewallConfig(
        security_groups={
            "icmp": SecurityGroup(
                port="any", proto=Protocol.ICMP, direction=Direction.ANY
            ),
            "out-any": SecurityGroup(
                port="any", proto=Protocol.ANY, direction=Direction.OUT
            ),
            "in-http": SecurityGroup(
                port=80, proto=Protocol.TCP, direction=Direction.IN
            ),
        },
        host_groups={"laptop": ["admin"]},
        firewall_rules={
            "laptop": {
                "all": ["icmp", "out-any"],
                "admin": ["in-http"],
            }
        },
        firewall_default={"all": ["icmp", "out-any"]},
    )


def test_resolve_host_specific_rules():
    cfg = _make_config()
    inbound, outbound = resolve_firewall_rules("laptop", cfg)

    assert {"port": "any", "proto": "icmp", "host": "any"} in inbound
    assert {"port": "any", "proto": "icmp", "host": "any"} in outbound
    assert {"port": "any", "proto": "any", "host": "any"} in outbound
    assert {"port": 80, "proto": "tcp", "group": "admin"} in inbound
    assert {"port": 80, "proto": "tcp", "group": "admin"} not in outbound


def test_resolve_fallback_to_default():
    cfg = _make_config()
    inbound, outbound = resolve_firewall_rules("unknown-host", cfg)

    assert {"port": "any", "proto": "icmp", "host": "any"} in inbound
    assert {"port": "any", "proto": "icmp", "host": "any"} in outbound
    assert {"port": "any", "proto": "any", "host": "any"} in outbound
    assert len(inbound) == 1
    assert len(outbound) == 2


def test_resolve_direction_any_goes_to_both():
    cfg = FirewallConfig(
        security_groups={
            "icmp": SecurityGroup(
                port="any", proto=Protocol.ICMP, direction=Direction.ANY
            ),
        },
        firewall_default={"all": ["icmp"]},
    )
    inbound, outbound = resolve_firewall_rules("any-host", cfg)
    assert len(inbound) == 1
    assert len(outbound) == 1
    assert inbound[0]["proto"] == "icmp"
    assert outbound[0]["proto"] == "icmp"


def test_resolve_peer_group_all_maps_to_host_any():
    cfg = FirewallConfig(
        security_groups={
            "out-any": SecurityGroup(
                port="any", proto=Protocol.ANY, direction=Direction.OUT
            ),
        },
        firewall_default={"all": ["out-any"]},
    )
    _, outbound = resolve_firewall_rules("host", cfg)
    assert outbound[0]["host"] == "any"
    assert "group" not in outbound[0]


def test_resolve_named_peer_group_maps_to_group():
    cfg = FirewallConfig(
        security_groups={
            "in-http": SecurityGroup(
                port=80, proto=Protocol.TCP, direction=Direction.IN
            ),
        },
        firewall_default={"admin": ["in-http"]},
    )
    inbound, _ = resolve_firewall_rules("host", cfg)
    assert inbound[0]["group"] == "admin"
    assert "host" not in inbound[0]
