"""
Power Save Analysis Module
Detects legacy power save, U-APSD, and TWT (Target Wake Time)
"""

from scapy.all import *
from scapy.layers.dot11 import *


def detect_power_save(packets):
    """Detect power save mechanisms (Legacy PS, U-APSD, TWT)"""
    power_save = {
        'legacy_ps': {'detected': False, 'ps_poll': 0, 'tim_broadcast': 0},
        'u_apsd': {'detected': False, 'count': 0, 'details': []},
        'twt': {'detected': False, 'agreements': 0, 'details': []},
        'power_mgmt_frames': 0,
        'clients_in_ps': set(),
        'efficiency': {'awake_time': 0, 'sleep_time': 0},
        'not_detected_reason': {}
    }

    for idx, pkt in enumerate(packets, 1):
        try:
            if pkt.haslayer(Dot11):
                dot11 = pkt[Dot11]

                # Power Management bit in Frame Control
                if dot11.FCfield & 0x10:  # Power Management bit
                    power_save['power_mgmt_frames'] += 1
                    if dot11.addr2:
                        power_save['clients_in_ps'].add(dot11.addr2)

                # PS-Poll frames (Legacy Power Save)
                if dot11.type == 1 and dot11.subtype == 0x0a:
                    power_save['legacy_ps']['detected'] = True
                    power_save['legacy_ps']['ps_poll'] += 1

            # TIM (Traffic Indication Map) in beacons
            if pkt.haslayer(Dot11Beacon):
                if pkt.haslayer(Dot11Elt):
                    elt = pkt[Dot11Elt]
                    while elt:
                        # TIM IE (ID 5)
                        if elt.ID == 5:
                            power_save['legacy_ps']['detected'] = True
                            power_save['legacy_ps']['tim_broadcast'] += 1

                        # U-APSD detection (WMM Info Element)
                        if elt.ID == 221 and len(elt.info) >= 7:
                            # WMM Info Element (OUI 00:50:F2, Type 2, Subtype 0)
                            if elt.info[0:3] == b'\x00\x50\xf2' and elt.info[3] == 0x02:
                                # U-APSD is indicated in QoS Info field
                                if len(elt.info) >= 8:
                                    qos_info = elt.info[7]
                                    # U-APSD enabled if bit 7 is set
                                    if qos_info & 0x80:
                                        power_save['u_apsd']['detected'] = True
                                        power_save['u_apsd']['count'] += 1
                                        if len(power_save['u_apsd']['details']) < 3:
                                            power_save['u_apsd']['details'].append({
                                                'packet': idx,
                                                'type': 'U-APSD Enabled in WMM'
                                            })

                        # TWT (Target Wake Time) IE - WiFi 6 feature (ID 255, Extension 216)
                        if elt.ID == 255 and len(elt.info) >= 2 and elt.info[0] == 216:
                            power_save['twt']['detected'] = True
                            power_save['twt']['agreements'] += 1
                            if len(power_save['twt']['details']) < 3:
                                power_save['twt']['details'].append({
                                    'packet': idx,
                                    'type': 'TWT Information Element'
                                })

                        if hasattr(elt, 'payload'):
                            elt = elt.payload.getlayer(Dot11Elt)
                        else:
                            break

            # TWT Setup Action frames (WiFi 6)
            if pkt.haslayer(Dot11) and dot11.type == 0 and dot11.subtype == 0x0d:
                if hasattr(pkt, 'load') and len(pkt.load) >= 2:
                    category = pkt.load[0]
                    # S1G (Sub-1 GHz) or HE action category might contain TWT
                    if category == 22:  # Protected HE
                        action = pkt.load[1]
                        if action in [6, 7]:  # TWT Setup/Teardown
                            power_save['twt']['detected'] = True
                            power_save['twt']['agreements'] += 1

        except Exception:
            continue

    # Convert set to list for JSON serialization
    power_save['clients_in_ps'] = list(power_save['clients_in_ps'])

    # Add explanations
    if not power_save['legacy_ps']['detected']:
        power_save['not_detected_reason']['legacy_ps'] = (
            "No legacy power save detected. Requires PS-Poll frames or TIM elements in beacons. "
            "Clients may not be using power save or capture missed these frames."
        )

    if not power_save['u_apsd']['detected']:
        power_save['not_detected_reason']['u_apsd'] = (
            "No U-APSD (Unscheduled Automatic Power Save Delivery) detected. "
            "Requires WMM Info Element with U-APSD bit set. U-APSD improves VoIP efficiency."
        )

    if not power_save['twt']['detected']:
        power_save['not_detected_reason']['twt'] = (
            "No TWT (Target Wake Time) detected. Requires WiFi 6 TWT IE (ID 255, Ext 216) "
            "or TWT Setup action frames. TWT schedules wake times for improved battery life."
        )

    return power_save
