import ipaddress
import logging
from typing import List, Dict, Any

from config import SERVER_PORTS, PRINTER_PORTS


class TopologyBuilder:
    def __init__(self):
        self.nodes: Dict[str, Dict[str, Any]] = {}
        self.edges: List[Dict[str, Any]] = []
        self.logger = logging.getLogger(__name__)

    def build_topology(self, scan_results: Dict[str, Any]) -> Dict[str, Any]:
        """
        Build network topology from scan results.

        Args:
            scan_results: Results from network scanner

        Returns:
            Dict with nodes and edges representing the topology
        """
        self.nodes.clear()
        self.edges.clear()
        hosts = scan_results.get('hosts', [])
        subnet = scan_results.get('subnet', '')

        if not hosts:
            return self.get_graph_data()

        for host in hosts:
            ip = host['ip']
            self.nodes[ip] = {
                'ip': ip,
                'hostname': host.get('hostname', ''),
                'state': host.get('state', 'unknown'),
                'mac_address': host.get('mac_address', ''),
                'vendor': host.get('vendor', ''),
                'open_ports': host.get('open_ports', []),
                'protocols': host.get('protocols', []),
                'node_type': self._classify_host(host),
            }

        self._infer_connections(hosts, subnet)

        return self.get_graph_data()

    def _classify_host(self, host: Dict[str, Any]) -> str:
        """Classify host type based on IP address and open ports."""
        ip = host['ip']
        if ip.endswith('.1') or ip.endswith('.254'):
            return 'gateway'

        host_ports = {port['port'] for port in host.get('open_ports', [])}

        if host_ports & SERVER_PORTS:
            return 'server'

        if host_ports & PRINTER_PORTS:
            return 'printer'

        return 'workstation'

    def _infer_connections(self, hosts: List[Dict[str, Any]], subnet: str):
        """
        Build a star topology: every host connects to the gateway.

        This is the only topology that can be honestly inferred from a
        scan without SNMP/CDP/LLDP neighbor data.
        """
        try:
            network = ipaddress.ip_network(subnet, strict=False)
            gateway_ip = str(network.network_address + 1)

            actual_gateway = None
            for ip, attrs in self.nodes.items():
                if attrs.get('node_type') == 'gateway':
                    actual_gateway = ip
                    break

            if actual_gateway:
                gateway_ip = actual_gateway

            if gateway_ip not in self.nodes:
                return

            host_ips = [h['ip'] for h in hosts]
            for ip in host_ips:
                if ip != gateway_ip:
                    self.edges.append({
                        'source': ip,
                        'target': gateway_ip,
                        'weight': 1.0,
                        'type': 'gateway',
                    })

        except (ValueError, ipaddress.AddressValueError) as exc:
            self.logger.error("Error inferring connections: %s", exc)

    def get_graph_data(self) -> Dict[str, Any]:
        """Get graph data in a format suitable for D3.js visualization."""
        nodes = []
        for ip, attrs in self.nodes.items():
            nodes.append({
                'id': ip,
                'label': attrs.get('hostname') or ip,
                'ip': ip,
                'type': attrs.get('node_type', 'workstation'),
                'hostname': attrs.get('hostname', ''),
                'state': attrs.get('state', 'unknown'),
                'mac_address': attrs.get('mac_address', ''),
                'vendor': attrs.get('vendor', ''),
                'open_ports': len(attrs.get('open_ports', [])),
                'protocols': attrs.get('protocols', []),
            })

        type_counts: Dict[str, int] = {}
        for attrs in self.nodes.values():
            nt = attrs.get('node_type', 'unknown')
            type_counts[nt] = type_counts.get(nt, 0) + 1

        return {
            'nodes': nodes,
            'links': list(self.edges),
            'stats': {
                'total_nodes': len(nodes),
                'total_links': len(self.edges),
                'node_types': type_counts,
            },
        }
