"""
QoS Analysis Module
Analyzes WMM/QoS, EDCA parameters, frame aggregation, and retry rates
"""

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


def detect_qos_analysis(packets):
    """Analyze QoS, WMM, EDCA parameters, and frame aggregation"""
    qos = {
        'wmm': {'detected': False, 'count': 0, 'ac_distribution': {'BE': 0, 'BK': 0, 'VI': 0, 'VO': 0}},
        'edca_params': {'detected': False, 'params': []},
        'tspec': {'detected': False, 'requests': 0, 'responses': 0},
        'frame_aggregation': {
            'a_mpdu': {'detected': False, 'count': 0},
            'a_msdu': {'detected': False, 'count': 0},
            'avg_aggregation': 0
        },
        'retry_analysis': {'total_retries': 0, 'retry_rate_pct': 0},
        'not_detected_reason': {}
    }

    total_qos_frames = 0
    retry_count = 0
    total_frames = 0

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

                # Retry detection
                if dot11.FCfield & 0x08:
                    retry_count += 1

                # QoS Data frames (type=2, subtype=8-15 with QoS bit)
                if dot11.type == 2 and (dot11.subtype & 0x08):
                    total_qos_frames += 1
                    qos['wmm']['detected'] = True
                    qos['wmm']['count'] += 1

                    # Extract TID (Traffic Identifier) from QoS Control field
                    if hasattr(pkt, 'load') and len(pkt) > 24:
                        # QoS Control is at offset 24-25 in QoS Data frame
                        # TID is bits 0-3 of QoS Control
                        # Access Category mapping: TID 0,3=BE, 1,2=BK, 4,5=VI, 6,7=VO
                        try:
                            qos_control_offset = 24
                            if hasattr(dot11, 'addr4'):  # 4-address frame
                                qos_control_offset = 30

                            # Try to extract TID (simplified, may need adjustment)
                            # In Scapy, QoS frames should have a QoS layer
                            if 'QoS' in str(type(pkt)):
                                tid = 0  # Default
                                # Map TID to AC
                                if tid in [1, 2]:
                                    qos['wmm']['ac_distribution']['BK'] += 1
                                elif tid in [0, 3]:
                                    qos['wmm']['ac_distribution']['BE'] += 1
                                elif tid in [4, 5]:
                                    qos['wmm']['ac_distribution']['VI'] += 1
                                elif tid in [6, 7]:
                                    qos['wmm']['ac_distribution']['VO'] += 1
                            else:
                                # Default to BE if can't determine
                                qos['wmm']['ac_distribution']['BE'] += 1
                        except:
                            qos['wmm']['ac_distribution']['BE'] += 1

            # A-MPDU detection (Aggregated MPDU in RadioTap or frame indicators)
            if pkt.haslayer(RadioTap):
                rt = pkt[RadioTap]
                # A-MPDU flag in RadioTap (if supported by capture)
                if hasattr(rt, 'A_MPDU_ref'):
                    qos['frame_aggregation']['a_mpdu']['detected'] = True
                    qos['frame_aggregation']['a_mpdu']['count'] += 1

            # EDCA Parameter Set IE (ID 12 in WMM Vendor Specific)
            if pkt.haslayer(Dot11Beacon) or pkt.haslayer(Dot11ProbeResp):
                if pkt.haslayer(Dot11Elt):
                    elt = pkt[Dot11Elt]
                    while elt:
                        # WMM Parameter Element (Vendor Specific, OUI 00:50:F2, Type 2)
                        if elt.ID == 221 and len(elt.info) >= 7:
                            if elt.info[0:3] == b'\x00\x50\xf2' and elt.info[3] == 0x02:
                                qos['edca_params']['detected'] = True
                                if len(qos['edca_params']['params']) < 3:
                                    qos['edca_params']['params'].append({
                                        'packet': idx,
                                        'type': 'WMM Parameter Element'
                                    })

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

            # TSPEC (Traffic Specification) - Action frames
            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]
                    if category == 17:  # WMM Action
                        qos['tspec']['detected'] = True
                        action = pkt.load[1]
                        if action == 0:  # ADDTS Request
                            qos['tspec']['requests'] += 1
                        elif action == 1:  # ADDTS Response
                            qos['tspec']['responses'] += 1

        except Exception:
            continue

    # Calculate retry rate
    if total_frames > 0:
        qos['retry_analysis']['total_retries'] = retry_count
        qos['retry_analysis']['retry_rate_pct'] = round((retry_count / total_frames) * 100, 2)

    # Add explanations
    if not qos['wmm']['detected']:
        qos['not_detected_reason']['wmm'] = (
            "No WMM/QoS detected. Requires QoS Data frames (type=2, QoS subtype bit set). "
            "WMM prioritizes traffic for voice, video, best effort, and background."
        )

    if not qos['edca_params']['detected']:
        qos['not_detected_reason']['edca'] = (
            "No EDCA parameters detected. Requires WMM Parameter Element in beacons/probe responses. "
            "EDCA defines contention windows and transmission opportunities per access category."
        )

    if not qos['frame_aggregation']['a_mpdu']['detected']:
        qos['not_detected_reason']['aggregation'] = (
            "No frame aggregation detected. A-MPDU/A-MSDU improve efficiency by combining multiple frames. "
            "May not be visible in RadioTap headers depending on capture device."
        )

    return qos
