from typing import Dict, List, Optional
from datetime import datetime, timedelta
import logging
import asyncio

logger = logging.getLogger(__name__)

class TrafficAnalyzer:
    """NetFlow/sFlow traffic analysis"""
    
    def __init__(self, db):
        self.db = db
        self.flow_collectors = {}
    
    async def process_netflow_packet(self, flow_data: Dict, org_id: str = None):
        """Process NetFlow packet"""
        try:
            flow_record = {
                'timestamp': datetime.utcnow(),
                'src_ip': flow_data.get('src_ip'),
                'dst_ip': flow_data.get('dst_ip'),
                'src_port': flow_data.get('src_port'),
                'dst_port': flow_data.get('dst_port'),
                'protocol': flow_data.get('protocol'),
                'bytes': flow_data.get('bytes', 0),
                'packets': flow_data.get('packets', 0),
                'flow_start': flow_data.get('flow_start'),
                'flow_end': flow_data.get('flow_end'),
                'exporter_ip': flow_data.get('exporter_ip'),
            }
            if org_id:
                flow_record['organization_id'] = org_id
            
            await self.db.netflow_records.insert_one(flow_record)
            
        except Exception as e:
            logger.error(f"Failed to process NetFlow packet: {e}")
    
    async def get_top_talkers(self, device_id: Optional[str] = None,
                             hours: int = 1, limit: int = 10,
                             org_id: str = None) -> List[Dict]:
        """Get top bandwidth consumers"""
        since = datetime.utcnow() - timedelta(hours=hours)
        
        match_stage = {'timestamp': {'$gte': since}}
        if org_id:
            match_stage['organization_id'] = org_id
        if device_id:
            device = await self.db.devices.find_one({'id': device_id})
            if device:
                match_stage['exporter_ip'] = device['ip']
        
        pipeline = [
            {'$match': match_stage},
            {'$group': {
                '_id': {'src_ip': '$src_ip', 'dst_ip': '$dst_ip'},
                'total_bytes': {'$sum': '$bytes'},
                'total_packets': {'$sum': '$packets'},
                'flow_count': {'$sum': 1}
            }},
            {'$sort': {'total_bytes': -1}},
            {'$limit': limit}
        ]
        
        results = await self.db.netflow_records.aggregate(pipeline).to_list(limit)
        
        return [
            {
                'src_ip': r['_id']['src_ip'],
                'dst_ip': r['_id']['dst_ip'],
                'total_bytes': r['total_bytes'],
                'total_packets': r['total_packets'],
                'flow_count': r['flow_count'],
                'avg_bytes_per_flow': r['total_bytes'] / r['flow_count'] if r['flow_count'] > 0 else 0
            }
            for r in results
        ]
    
    async def get_top_applications(self, hours: int = 1, limit: int = 10,
                                   org_id: str = None) -> List[Dict]:
        """Get top applications by port"""
        since = datetime.utcnow() - timedelta(hours=hours)
        
        match_filter = {'timestamp': {'$gte': since}}
        if org_id:
            match_filter['organization_id'] = org_id
        
        pipeline = [
            {'$match': match_filter},
            {'$group': {
                '_id': '$dst_port',
                'total_bytes': {'$sum': '$bytes'},
                'total_packets': {'$sum': '$packets'},
                'flow_count': {'$sum': 1}
            }},
            {'$sort': {'total_bytes': -1}},
            {'$limit': limit}
        ]
        
        results = await self.db.netflow_records.aggregate(pipeline).to_list(limit)
        
        return [
            {
                'port': r['_id'],
                'application': self._port_to_application(r['_id']),
                'total_bytes': r['total_bytes'],
                'total_packets': r['total_packets'],
                'flow_count': r['flow_count']
            }
            for r in results
        ]
    
    async def get_top_protocols(self, hours: int = 1, org_id: str = None) -> List[Dict]:
        """Get traffic breakdown by protocol"""
        since = datetime.utcnow() - timedelta(hours=hours)
        
        match_filter = {'timestamp': {'$gte': since}}
        if org_id:
            match_filter['organization_id'] = org_id
        
        pipeline = [
            {'$match': match_filter},
            {'$group': {
                '_id': '$protocol',
                'total_bytes': {'$sum': '$bytes'},
                'total_packets': {'$sum': '$packets'},
                'flow_count': {'$sum': 1}
            }},
            {'$sort': {'total_bytes': -1}}
        ]
        
        results = await self.db.netflow_records.aggregate(pipeline).to_list(100)
        
        return [
            {
                'protocol': self._protocol_number_to_name(r['_id']),
                'protocol_number': r['_id'],
                'total_bytes': r['total_bytes'],
                'total_packets': r['total_packets'],
                'flow_count': r['flow_count']
            }
            for r in results
        ]
    
    async def get_traffic_timeline(self, hours: int = 24, interval_minutes: int = 30,
                                   org_id: str = None) -> List[Dict]:
        """Get traffic over time"""
        since = datetime.utcnow() - timedelta(hours=hours)
        
        match_filter = {'timestamp': {'$gte': since}}
        if org_id:
            match_filter['organization_id'] = org_id
        
        pipeline = [
            {'$match': match_filter},
            {'$group': {
                '_id': {
                    'year': {'$year': '$timestamp'},
                    'month': {'$month': '$timestamp'},
                    'day': {'$dayOfMonth': '$timestamp'},
                    'hour': {'$hour': '$timestamp'},
                    'minute': {'$subtract': [
                        {'$minute': '$timestamp'},
                        {'$mod': [{'$minute': '$timestamp'}, interval_minutes]}
                    ]}
                },
                'total_bytes': {'$sum': '$bytes'},
                'total_packets': {'$sum': '$packets'},
                'flow_count': {'$sum': 1}
            }},
            {'$sort': {'_id': 1}}
        ]
        
        results = await self.db.netflow_records.aggregate(pipeline).to_list(1000)
        
        return [
            {
                'timestamp': datetime(
                    r['_id']['year'],
                    r['_id']['month'],
                    r['_id']['day'],
                    r['_id']['hour'],
                    r['_id']['minute']
                ).isoformat(),
                'total_bytes': r['total_bytes'],
                'total_packets': r['total_packets'],
                'flow_count': r['flow_count'],
                'bps': (r['total_bytes'] * 8) / (interval_minutes * 60)  # bits per second
            }
            for r in results
        ]
    
    async def get_conversation_pairs(self, hours: int = 1, limit: int = 20,
                                     org_id: str = None) -> List[Dict]:
        """Get top conversation pairs"""
        since = datetime.utcnow() - timedelta(hours=hours)
        
        match_filter = {'timestamp': {'$gte': since}}
        if org_id:
            match_filter['organization_id'] = org_id
        
        pipeline = [
            {'$match': match_filter},
            {'$group': {
                '_id': {
                    'ip1': {'$min': ['$src_ip', '$dst_ip']},
                    'ip2': {'$max': ['$src_ip', '$dst_ip']}
                },
                'total_bytes': {'$sum': '$bytes'},
                'total_packets': {'$sum': '$packets'}
            }},
            {'$sort': {'total_bytes': -1}},
            {'$limit': limit}
        ]
        
        results = await self.db.netflow_records.aggregate(pipeline).to_list(limit)
        
        return [
            {
                'host1': r['_id']['ip1'],
                'host2': r['_id']['ip2'],
                'total_bytes': r['total_bytes'],
                'total_packets': r['total_packets']
            }
            for r in results
        ]
    
    async def analyze_bandwidth_usage(self, device_id: str, hours: int = 24,
                                      org_id: str = None) -> Dict:
        """Comprehensive bandwidth analysis for device"""
        query = {'id': device_id}
        if org_id:
            query['organization_id'] = org_id
        device = await self.db.devices.find_one(query)
        if not device:
            raise Exception("Device not found")
        
        since = datetime.utcnow() - timedelta(hours=hours)
        
        # Total traffic
        pipeline = [
            {'$match': {
                'exporter_ip': device['ip'],
                'timestamp': {'$gte': since}
            }},
            {'$group': {
                '_id': None,
                'total_bytes': {'$sum': '$bytes'},
                'total_packets': {'$sum': '$packets'},
                'flow_count': {'$sum': 1}
            }}
        ]
        
        total_result = await self.db.netflow_records.aggregate(pipeline).to_list(1)
        total_stats = total_result[0] if total_result else {
            'total_bytes': 0,
            'total_packets': 0,
            'flow_count': 0
        }
        
        top_talkers = await self.get_top_talkers(device_id, hours, 5, org_id=org_id)
        top_apps = await self.get_top_applications(hours, 5, org_id=org_id)
        top_protocols = await self.get_top_protocols(hours, org_id=org_id)
        
        return {
            'device_id': device_id,
            'device_ip': device['ip'],
            'period_hours': hours,
            'total_stats': total_stats,
            'average_bps': (total_stats['total_bytes'] * 8) / (hours * 3600) if hours > 0 else 0,
            'top_talkers': top_talkers,
            'top_applications': top_apps,
            'top_protocols': top_protocols
        }
    
    def _port_to_application(self, port: int) -> str:
        """Map port number to application name"""
        port_map = {
            20: 'FTP-Data',
            21: 'FTP',
            22: 'SSH',
            23: 'Telnet',
            25: 'SMTP',
            53: 'DNS',
            80: 'HTTP',
            110: 'POP3',
            143: 'IMAP',
            443: 'HTTPS',
            3306: 'MySQL',
            3389: 'RDP',
            5432: 'PostgreSQL',
            8080: 'HTTP-Alt',
            8443: 'HTTPS-Alt',
        }
        return port_map.get(port, f'Port-{port}')
    
    def _protocol_number_to_name(self, proto_num: int) -> str:
        """Map protocol number to name"""
        proto_map = {
            1: 'ICMP',
            6: 'TCP',
            17: 'UDP',
            47: 'GRE',
            50: 'ESP',
            51: 'AH',
            89: 'OSPF',
        }
        return proto_map.get(proto_num, f'Protocol-{proto_num}')
    
    async def cleanup_old_flows(self, days_to_keep: int = 7, org_id: str = None):
        """Clean up old flow records, optionally scoped to organization"""
        cutoff_date = datetime.utcnow() - timedelta(days=days_to_keep)
        
        query = {'timestamp': {'$lt': cutoff_date}}
        if org_id:
            query['organization_id'] = org_id
        
        result = await self.db.netflow_records.delete_many(query)
        
        logger.info(f"Deleted {result.deleted_count} old flow records")
        return result.deleted_count