from collections import defaultdict
from enrichment.utils import get_channel_from_freq

def get_clients_on_ap(capture, ap_bssid):
    clients = defaultdict(int)
    ap_bssid = ap_bssid.lower()

    for packet in capture:
        try:
            if not hasattr(packet, "wlan"):
                continue

            sa = getattr(packet.wlan, "sa", '').lower()
            da = getattr(packet.wlan, "da", '').lower()
            bssid = getattr(packet.wlan, "bssid", '').lower()

            if bssid == ap_bssid or sa == ap_bssid or da == ap_bssid:
                if sa == ap_bssid and da and da != ap_bssid and not da.startswith("ff:ff:ff:ff:ff:ff"):
                    clients[da] += 1
                elif sa and sa != ap_bssid and not sa.startswith("ff:ff:ff:ff:ff:ff"):
                    clients[sa] += 1
        except AttributeError:
            continue

    return len([mac for mac, count in clients.items() if count > 3])

def get_clients_on_channel(capture, ap_channel, ap_bssid):
    try:
        ap_channel = int(ap_channel)
    except ValueError:
        print(f"[!] Could not parse channel number: {ap_channel}")
        return 0

    clients = set()

    for packet in capture:
        try:
            if 'radiotap' not in packet or 'wlan' not in packet:
                continue

            radio = packet.radiotap
            wlan = packet.wlan

            if not hasattr(radio, 'channel') or not hasattr(radio.channel, 'freq'):
                continue

            packet_freq = int(radio.channel.freq)
            packet_channel = get_channel_from_freq(packet_freq)

            if packet_channel != ap_channel:
                continue

            sa = getattr(wlan, 'sa', '').lower()
            da = getattr(wlan, 'da', '').lower()

            for mac in (sa, da):
                if mac and mac != 'ff:ff:ff:ff:ff:ff' and mac != ap_bssid:
                    clients.add(mac)

        except AttributeError:
            continue
        except Exception as e:
            print(f"[!] Error parsing packet: {e}")
            continue

    return len(clients)