from graph.topology_builder import TopologyBuilder


SAMPLE_SCAN = {
    'subnet': '192.168.1.0/24',
    'hosts': [
        {
            'ip': '192.168.1.1',
            'hostname': 'gateway',
            'state': 'up',
            'mac_address': 'AA:BB:CC:DD:EE:01',
            'vendor': 'Cisco',
            'open_ports': [],
            'protocols': [],
        },
        {
            'ip': '192.168.1.10',
            'hostname': 'web-server',
            'state': 'up',
            'mac_address': 'AA:BB:CC:DD:EE:10',
            'vendor': 'Dell',
            'open_ports': [
                {'port': 80, 'protocol': 'tcp', 'service': 'http'},
                {'port': 443, 'protocol': 'tcp', 'service': 'https'},
            ],
            'protocols': ['tcp'],
        },
        {
            'ip': '192.168.1.50',
            'hostname': '',
            'state': 'up',
            'mac_address': 'AA:BB:CC:DD:EE:50',
            'vendor': '',
            'open_ports': [],
            'protocols': [],
        },
        {
            'ip': '192.168.1.99',
            'hostname': 'office-printer',
            'state': 'up',
            'mac_address': 'AA:BB:CC:DD:EE:99',
            'vendor': 'HP',
            'open_ports': [
                {'port': 9100, 'protocol': 'tcp', 'service': 'jetdirect'},
            ],
            'protocols': ['tcp'],
        },
    ],
}


class TestClassifyHost:
    def setup_method(self):
        self.builder = TopologyBuilder()

    def test_gateway_dot_1(self):
        host = {'ip': '192.168.1.1', 'open_ports': []}
        assert self.builder._classify_host(host) == 'gateway'

    def test_gateway_dot_254(self):
        host = {'ip': '10.0.0.254', 'open_ports': []}
        assert self.builder._classify_host(host) == 'gateway'

    def test_server_with_http(self):
        host = {'ip': '192.168.1.10', 'open_ports': [{'port': 80}]}
        assert self.builder._classify_host(host) == 'server'

    def test_server_with_ssh(self):
        host = {'ip': '192.168.1.10', 'open_ports': [{'port': 22}]}
        assert self.builder._classify_host(host) == 'server'

    def test_printer(self):
        host = {'ip': '192.168.1.99', 'open_ports': [{'port': 9100}]}
        assert self.builder._classify_host(host) == 'printer'

    def test_workstation_default(self):
        host = {'ip': '192.168.1.50', 'open_ports': []}
        assert self.builder._classify_host(host) == 'workstation'

    def test_gateway_takes_priority_over_server(self):
        host = {'ip': '192.168.1.1', 'open_ports': [{'port': 80}]}
        assert self.builder._classify_host(host) == 'gateway'


class TestBuildTopology:
    def test_basic_build(self):
        builder = TopologyBuilder()
        result = builder.build_topology(SAMPLE_SCAN)

        assert result['stats']['total_nodes'] == 4
        assert len(result['nodes']) == 4
        assert len(result['links']) == 3

    def test_all_edges_connect_to_gateway(self):
        builder = TopologyBuilder()
        result = builder.build_topology(SAMPLE_SCAN)

        for link in result['links']:
            assert '192.168.1.1' in (link['source'], link['target'])

    def test_edge_type_is_gateway(self):
        builder = TopologyBuilder()
        result = builder.build_topology(SAMPLE_SCAN)

        for link in result['links']:
            assert link['type'] == 'gateway'

    def test_empty_hosts(self):
        builder = TopologyBuilder()
        result = builder.build_topology({'subnet': '10.0.0.0/24', 'hosts': []})

        assert result['stats']['total_nodes'] == 0
        assert result['links'] == []

    def test_no_gateway_in_hosts(self):
        builder = TopologyBuilder()
        scan = {
            'subnet': '192.168.1.0/24',
            'hosts': [
                {'ip': '192.168.1.50', 'open_ports': [], 'state': 'up'},
                {'ip': '192.168.1.60', 'open_ports': [], 'state': 'up'},
            ],
        }
        result = builder.build_topology(scan)

        assert result['stats']['total_nodes'] == 2
        assert result['links'] == []


class TestGetGraphData:
    def test_output_shape(self):
        builder = TopologyBuilder()
        builder.build_topology(SAMPLE_SCAN)
        data = builder.get_graph_data()

        assert 'nodes' in data
        assert 'links' in data
        assert 'stats' in data
        assert 'total_nodes' in data['stats']
        assert 'total_links' in data['stats']
        assert 'node_types' in data['stats']

    def test_node_fields(self):
        builder = TopologyBuilder()
        builder.build_topology(SAMPLE_SCAN)
        data = builder.get_graph_data()

        node = data['nodes'][0]
        for field in ('id', 'label', 'ip', 'type', 'hostname', 'state',
                      'mac_address', 'vendor', 'open_ports', 'protocols'):
            assert field in node

    def test_node_type_counts(self):
        builder = TopologyBuilder()
        builder.build_topology(SAMPLE_SCAN)
        data = builder.get_graph_data()

        counts = data['stats']['node_types']
        assert counts.get('gateway') == 1
        assert counts.get('server') == 1
        assert counts.get('workstation') == 1
        assert counts.get('printer') == 1
