"""
PCAP file loading utilities
"""

import sys
from typing import List, Iterator, Optional

try:
    from scapy.all import rdpcap, PcapReader, Packet
except ImportError:
    print("Error: scapy library required. Install with: pip install scapy")
    sys.exit(1)


class PCAPLoader:
    """Utility class for loading PCAP files"""
    
    def __init__(self, file_path: str):
        self.file_path = file_path
        self._packet_count: Optional[int] = None
    
    def load_all(self) -> List[Packet]:
        """Load all packets from the PCAP file"""
        try:
            packets = rdpcap(self.file_path)
            self._packet_count = len(packets)
            return packets
        except Exception as e:
            raise IOError(f"Error loading PCAP file {self.file_path}: {e}")
    
    def load_streaming(self, chunk_size: int = 1000) -> Iterator[List[Packet]]:
        """Load packets in chunks for memory efficiency"""
        try:
            with PcapReader(self.file_path) as pcap_reader:
                chunk = []
                for packet in pcap_reader:
                    chunk.append(packet)
                    if len(chunk) >= chunk_size:
                        yield chunk
                        chunk = []
                
                # Yield remaining packets
                if chunk:
                    yield chunk
                    
        except Exception as e:
            raise IOError(f"Error streaming PCAP file {self.file_path}: {e}")
    
    def get_packet_count(self) -> Optional[int]:
        """Get the total number of packets (if loaded)"""
        return self._packet_count
    
    def validate_file(self) -> bool:
        """Validate that the file exists and is a valid PCAP"""
        try:
            with PcapReader(self.file_path) as pcap_reader:
                # Try to read first packet
                next(iter(pcap_reader))
                return True
        except (IOError, StopIteration):
            return False
        except Exception:
            return False
    
    @staticmethod
    def get_file_info(file_path: str) -> dict:
        """Get basic information about a PCAP file"""
        try:
            packet_count = 0
            first_timestamp = None
            last_timestamp = None
            total_bytes = 0
            
            with PcapReader(file_path) as pcap_reader:
                for packet in pcap_reader:
                    packet_count += 1
                    total_bytes += len(packet)
                    
                    if first_timestamp is None:
                        first_timestamp = float(packet.time)
                    last_timestamp = float(packet.time)
            
            duration = (last_timestamp - first_timestamp) if first_timestamp and last_timestamp else 0
            
            return {
                'file_path': file_path,
                'packet_count': packet_count,
                'total_bytes': total_bytes,
                'duration_seconds': duration,
                'first_timestamp': first_timestamp,
                'last_timestamp': last_timestamp,
                'avg_packet_rate': packet_count / duration if duration > 0 else 0
            }
            
        except Exception as e:
            return {'error': str(e)}