#!/usr/bin/env python3

import argparse

from scapy.all import get_if_list

from frametypes import FRAME_TYPES
from core.analyzer import PacketFlowAnalyzer


def main():
    p = argparse.ArgumentParser(description="Airstream - Packet Flow Analyzer")
    p.add_argument('-p', '--pcap', help='PCAP file')
    p.add_argument('-i', '--interface', help='Network interface')
    p.add_argument('-c', '--count', type=int, default=100, help='Packet count (default: 100)')
    p.add_argument('-l', '--list-interfaces', action='store_true', help='List interfaces')
    p.add_argument('-o', '--output', help='CSV output file')
    p.add_argument('-s', '--stats', nargs='+', choices=list(FRAME_TYPES.keys()) + ['all'], 
                   default=['overview'], help='Statistics types to use (default: overview, use "all" for all types)')
    
    args = p.parse_args()
    
    if args.list_interfaces:
        print("Interfaces:", *get_if_list(), sep='\n  ')
        return
    
    if not (args.pcap or args.interface):
        p.error("Specify --pcap or --interface")
    
    # Handle 'all' option
    if 'all' in args.stats:
        selected_stats = list(FRAME_TYPES.keys())
        stats_classes = list(FRAME_TYPES.values())
    else:
        selected_stats = args.stats
        stats_classes = [FRAME_TYPES[s] for s in args.stats]
    
    print(f"Using stats: {', '.join(selected_stats)}")
    
    analyzer = PacketFlowAnalyzer(stats_classes)
    
    try:
        if args.pcap:
            analyzer.analyze_pcap(args.pcap)
        else:
            analyzer.analyze_live(args.interface, args.count)
        
        analyzer.print_summary()
        
        if args.output:
            analyzer.summary().to_csv(args.output, index=False)
            print(f"Saved: {args.output}")
            
    except KeyboardInterrupt:
        print("\nInterrupted")
    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()