import asyncio
import logging
from typing import Dict, List, Optional, Tuple
from pysnmp.hlapi.asyncio import *

logger = logging.getLogger(__name__)

class LLDPCDPDiscovery:
    """LLDP and CDP Discovery Module"""
    
    # LLDP OIDs
    LLDP_REM_CHASSIS_ID = '1.0.8802.1.1.2.1.4.1.1.5'
    LLDP_REM_PORT_ID = '1.0.8802.1.1.2.1.4.1.1.7'
    LLDP_REM_PORT_DESC = '1.0.8802.1.1.2.1.4.1.1.8'
    LLDP_REM_SYS_NAME = '1.0.8802.1.1.2.1.4.1.1.9'
    LLDP_REM_SYS_DESC = '1.0.8802.1.1.2.1.4.1.1.10'
    LLDP_REM_SYS_CAP = '1.0.8802.1.1.2.1.4.1.1.12'
    
    # CDP OIDs (Cisco)
    CDP_CACHE_DEVICE_ID = '1.3.6.1.4.1.9.9.23.1.2.1.1.6'
    CDP_CACHE_ADDRESS = '1.3.6.1.4.1.9.9.23.1.2.1.1.4'
    CDP_CACHE_VERSION = '1.3.6.1.4.1.9.9.23.1.2.1.1.5'
    CDP_CACHE_PLATFORM = '1.3.6.1.4.1.9.9.23.1.2.1.1.8'
    CDP_CACHE_PORT_ID = '1.3.6.1.4.1.9.9.23.1.2.1.1.7'
    
    async def discover_lldp_neighbors(self, ip: str, community: str = 'public') -> List[Dict]:
        """Discover LLDP neighbors"""
        neighbors = []
        
        try:
            # Walk LLDP remote table
            async for (errorIndication, errorStatus, errorIndex, varBinds) in nextCmd(
                SnmpEngine(),
                CommunityData(community),
                await UdpTransportTarget.create((ip, 161), timeout=5, retries=1),
                ContextData(),
                ObjectType(ObjectIdentity(self.LLDP_REM_SYS_NAME)),
                lexicographicMode=False
            ):
                if errorIndication or errorStatus:
                    break
                
                for varBind in varBinds:
                    oid = str(varBind[0])
                    value = str(varBind[1])
                    
                    # Extract index from OID
                    index = oid.split('.')[-3:]  # time, local_port, index
                    index_str = '.'.join(index)
                    
                    # Get additional info for this neighbor
                    neighbor_info = await self._get_lldp_neighbor_info(
                        ip, community, index_str
                    )
                    
                    if neighbor_info:
                        neighbors.append(neighbor_info)
            
            logger.info(f"Found {len(neighbors)} LLDP neighbors on {ip}")
            
        except Exception as e:
            logger.error(f"LLDP discovery failed for {ip}: {e}")
        
        return neighbors
    
    async def _get_lldp_neighbor_info(self, ip: str, community: str, index: str) -> Optional[Dict]:
        """Get detailed info for LLDP neighbor"""
        try:
            oids = [
                f"{self.LLDP_REM_CHASSIS_ID}.{index}",
                f"{self.LLDP_REM_PORT_ID}.{index}",
                f"{self.LLDP_REM_PORT_DESC}.{index}",
                f"{self.LLDP_REM_SYS_NAME}.{index}",
                f"{self.LLDP_REM_SYS_DESC}.{index}",
            ]
            
            results = {}
            for oid in oids:
                errorIndication, errorStatus, errorIndex, varBinds = await getCmd(
                    SnmpEngine(),
                    CommunityData(community),
                    await UdpTransportTarget.create((ip, 161), timeout=2, retries=0),
                    ContextData(),
                    ObjectType(ObjectIdentity(oid))
                )
                
                if not errorIndication and not errorStatus:
                    for varBind in varBinds:
                        results[oid] = str(varBind[1])
            
            neighbor = {
                'protocol': 'LLDP',
                'chassis_id': results.get(oids[0], 'N/A'),
                'port_id': results.get(oids[1], 'N/A'),
                'port_description': results.get(oids[2], 'N/A'),
                'system_name': results.get(oids[3], 'N/A'),
                'system_description': results.get(oids[4], 'N/A'),
            }
            
            return neighbor
            
        except Exception as e:
            logger.error(f"Failed to get LLDP neighbor info: {e}")
            return None
    
    async def discover_cdp_neighbors(self, ip: str, community: str = 'public') -> List[Dict]:
        """Discover CDP neighbors (Cisco)"""
        neighbors = []
        
        try:
            # Walk CDP cache table
            async for (errorIndication, errorStatus, errorIndex, varBinds) in nextCmd(
                SnmpEngine(),
                CommunityData(community),
                await UdpTransportTarget.create((ip, 161), timeout=5, retries=1),
                ContextData(),
                ObjectType(ObjectIdentity(self.CDP_CACHE_DEVICE_ID)),
                lexicographicMode=False
            ):
                if errorIndication or errorStatus:
                    break
                
                for varBind in varBinds:
                    oid = str(varBind[0])
                    device_id = str(varBind[1])
                    
                    # Extract index from OID
                    index_parts = oid.split('.')[-2:]
                    index_str = '.'.join(index_parts)
                    
                    # Get additional info
                    neighbor_info = await self._get_cdp_neighbor_info(
                        ip, community, index_str, device_id
                    )
                    
                    if neighbor_info:
                        neighbors.append(neighbor_info)
            
            logger.info(f"Found {len(neighbors)} CDP neighbors on {ip}")
            
        except Exception as e:
            logger.error(f"CDP discovery failed for {ip}: {e}")
        
        return neighbors
    
    async def _get_cdp_neighbor_info(self, ip: str, community: str, index: str, device_id: str) -> Optional[Dict]:
        """Get detailed info for CDP neighbor"""
        try:
            oids = [
                f"{self.CDP_CACHE_PORT_ID}.{index}",
                f"{self.CDP_CACHE_PLATFORM}.{index}",
                f"{self.CDP_CACHE_VERSION}.{index}",
            ]
            
            results = {}
            for oid in oids:
                errorIndication, errorStatus, errorIndex, varBinds = await getCmd(
                    SnmpEngine(),
                    CommunityData(community),
                    await UdpTransportTarget.create((ip, 161), timeout=2, retries=0),
                    ContextData(),
                    ObjectType(ObjectIdentity(oid))
                )
                
                if not errorIndication and not errorStatus:
                    for varBind in varBinds:
                        results[oid] = str(varBind[1])
            
            neighbor = {
                'protocol': 'CDP',
                'device_id': device_id,
                'port_id': results.get(oids[0], 'N/A'),
                'platform': results.get(oids[1], 'N/A'),
                'version': results.get(oids[2], 'N/A'),
            }
            
            return neighbor
            
        except Exception as e:
            logger.error(f"Failed to get CDP neighbor info: {e}")
            return None
    
    async def discover_all_neighbors(self, ip: str, community: str = 'public') -> Dict:
        """Discover both LLDP and CDP neighbors"""
        lldp_neighbors, cdp_neighbors = await asyncio.gather(
            self.discover_lldp_neighbors(ip, community),
            self.discover_cdp_neighbors(ip, community),
            return_exceptions=True
        )
        
        if isinstance(lldp_neighbors, Exception):
            lldp_neighbors = []
        if isinstance(cdp_neighbors, Exception):
            cdp_neighbors = []
        
        return {
            'lldp': lldp_neighbors,
            'cdp': cdp_neighbors,
            'total': len(lldp_neighbors) + len(cdp_neighbors)
        }

# Global instance
lldp_cdp_discovery = LLDPCDPDiscovery()