from collections import defaultdict
from typing import Optional, List, Type

import pandas as pd
from scapy.all import PcapReader, IP, TCP, UDP, sniff
from tabulate import tabulate

from frametypes import FrameTypeInterface, FRAME_TYPES
from core.models import FlowKey
from core.stats import MultiStats


class PacketFlowAnalyzer:
    def __init__(self, stats_classes: List[Type[FrameTypeInterface]] = None):
        if stats_classes is None:
            stats_classes = [FRAME_TYPES['overview']]
        self.stats_classes = stats_classes
        self.flows = defaultdict(lambda: MultiStats(stats_classes))
    
    def _get_flow_key(self, pkt) -> Optional[FlowKey]:
        if not pkt.haslayer(IP):
            return None
        
        ip = pkt[IP]
        if pkt.haslayer(TCP):
            return FlowKey(ip.src, pkt[TCP].sport, ip.dst, pkt[TCP].dport, "TCP")
        elif pkt.haslayer(UDP):
            return FlowKey(ip.src, pkt[UDP].sport, ip.dst, pkt[UDP].dport, "UDP")
        else:
            return FlowKey(ip.src, 0, ip.dst, 0, str(ip.proto))
    
    def _process(self, pkt):
        key = self._get_flow_key(pkt)
        if key:
            self.flows[key].add(float(pkt.time), len(pkt), pkt)
    
    def analyze_pcap(self, file: str):
        print(f"Analyzing: {file}")
        for pkt in PcapReader(file):
            self._process(pkt)
        print(f"Found {len(self.flows)} flows")
    
    def analyze_live(self, iface: str, count: int = 100):
        print(f"Capturing {count} packets on {iface}")
        sniff(iface=iface, prn=self._process, count=count)
        print(f"Found {len(self.flows)} flows")
    
    def summary(self) -> pd.DataFrame:
        rows = []
        for k, multi_stats in self.flows.items():
            row = {
                'Src IP': k.src_ip,
                'Src Port': k.src_port,
                'Dst IP': k.dst_ip,
                'Dst Port': k.dst_port,
                'Proto': k.protocol
            }
            if k.extended_type:
                row['Type'] = k.extended_type
            row.update(multi_stats.get_combined_summary())
            rows.append(row)
        return pd.DataFrame(rows)
    
    def print_summary(self):
        df = self.summary()
        if df.empty:
            print("No flows detected")
            return
        
        print(f"\n{len(df)} flows:")
        print(tabulate(df, headers='keys', tablefmt='plain', showindex=False))
        print(f"\nTotals: {df['Pkts'].sum()} packets, {df['Bytes'].sum()} bytes")