import psutil
import os, time
import socket, fcntl, struct
import whatap.util.process_util as process_util
from whatap.util.hash_util import HashUtil as hash
from whatap.util.ip_util import IPUtil as ip_util
from whatap.pack.pack_enum import PackEnum
import whatap.pack.smnetperf_pack as smnetperf_pack
from StringIO import StringIO

class _NetPerfRaw:
    def __init__(self):
        self.deviceID = None
        self.readCount = None
        self.readByteCount = None
        self.readDroppedCount = None
        self.readErrorCount = None
        self.writeCount = None
        self.writeByteCount = None
        self.writeDroppedCount = None
        self.writeErrorCount = None
        self.timestamp = None

class _NetPerf:
    def __init__(self):
        self.desc = None
        self.ip = None
        self.hwAddr = None
        self.trafficIn = None
        self.trafficOut = None
        self.packetIn = None
        self.packetOut = None
        self.errorIn = None
        self.errorOut =  None
        self.droppedOut = None
        self.droppedIn = None
        self.linkUsageMax = None
        self.nicName = None

    def __str__(self):
        buf = StringIO()
        buf.write('\ndesc=')
        buf.write(self.desc)
        buf.write('\nip=')
        buf.write(self.ip)
        buf.write('\nhwAddr=')
        buf.write(self.hwAddr)
        buf.write('\ntrafficIn=')
        buf.write(self.trafficIn)
        buf.write('\ntrafficOut=')
        buf.write(self.trafficOut)
        buf.write('\npacketIn=')
        buf.write(self.packetIn)
        buf.write('\npacketOut=')
        buf.write(self.packetOut)
        buf.write('\nerrorIn=')
        buf.write(str(self.errorIn))
        buf.write('\nerrorOut=')
        buf.write(str(self.errorOut))
        buf.write('\ndroppedOut=')
        buf.write(str(self.droppedOut))
        buf.write('\ndroppedIn=')
        buf.write(str(self.droppedIn))

        return buf.getvalue()

_nicInfoDict = {}

def getNicInfo(deviceId):
    
    if deviceId not in _nicInfoDict:
        addrs = psutil.net_if_addrs()
        if deviceId not in addrs or not addrs[deviceId]:
            return None, None
        nicip = addrs[deviceId][0].address
        hwaddr = None
        entstat = process_util.executeCommandShellWithTimeout('entstat -d %s' % deviceId, timeout=10)
        for l in entstat.split('\n'):
            l = l.strip()
            if l and l.startswith('Hardware Address:'):
                hwaddr = l.split(':')[1].strip()
                break
        if not hwaddr:
            return None, None
        _nicInfoDict[deviceId] = (nicip, hwaddr)
        
    nicip, hwaddr = _nicInfoDict[deviceId]

    return nicip, hwaddr

oldNicTrafficMap = None
def GetNicUtil():
    global oldNicTrafficMap
    iocounters = psutil.net_io_counters(pernic=True)
    
    nicTrafficMap = {}

    timestamp = int(time.time() * 1000)
    traffics = []
    j =0
    for deviceId, iocounter in iocounters.items():
        nicTrafficIoRaw = _NetPerfRaw()
        nicTrafficIoRaw.deviceId = deviceId
        nicTrafficIoRaw.readByteCount = int(iocounter.bytes_recv)
        nicTrafficIoRaw.readCount = int(iocounter.packets_recv)
        nicTrafficIoRaw.readErrorCount = int(iocounter.errin)
        nicTrafficIoRaw.readDroppedCount = int(iocounter.dropin)
        nicTrafficIoRaw.writeByteCount = int(iocounter.bytes_sent)
        nicTrafficIoRaw.writeCount = int(iocounter.packets_sent)
        nicTrafficIoRaw.writeErrorCount = int(iocounter.errout)
        nicTrafficIoRaw.writeDroppedCount = int(iocounter.dropout)
        nicTrafficIoRaw.timestamp = timestamp

        nicTrafficMap[deviceId] = nicTrafficIoRaw 
        if oldNicTrafficMap and deviceId in oldNicTrafficMap:
            oldNicTrafficIoRaw = oldNicTrafficMap[deviceId]
            timediff = nicTrafficIoRaw.timestamp - oldNicTrafficIoRaw.timestamp
            n = _NetPerf()
            n.desc = deviceId
            n.trafficIn = float(nicTrafficIoRaw.readByteCount - oldNicTrafficIoRaw.readByteCount) / timediff * 1000.0 * 8
            n.trafficOut = float(nicTrafficIoRaw.writeByteCount - oldNicTrafficIoRaw.writeByteCount) / timediff * 1000.0 * 8
            n.packetIn = float(nicTrafficIoRaw.readCount - oldNicTrafficIoRaw.readCount) / timediff * 1000.0
            n.packetOut = float(nicTrafficIoRaw.writeCount - oldNicTrafficIoRaw.writeCount) / timediff * 1000.0
            n.droppedIn = float(nicTrafficIoRaw.readDroppedCount - oldNicTrafficIoRaw.readDroppedCount) / timediff * 1000.0
            n.droppedOut = float(nicTrafficIoRaw.writeDroppedCount - oldNicTrafficIoRaw.writeDroppedCount) / timediff * 1000.0
            n.errorIn = float(nicTrafficIoRaw.readErrorCount - oldNicTrafficIoRaw.readErrorCount) / timediff * 1000.0
            n.errorOut = float(nicTrafficIoRaw.writeErrorCount - oldNicTrafficIoRaw.writeErrorCount) / timediff * 1000.0
            nicip, hwaddr = getNicInfo(deviceId)
            if not nicip or not hwaddr:
                continue
            n.ip = ip_util.toBytes(nicip)
            n.hwAddr = hwaddr

            traffics.append(n)
    oldNicTrafficMap = nicTrafficMap

    return traffics


_nicLinkSpeedCache = {}

def GetNicLinkSpeed(iface):
    if iface in _nicLinkSpeedCache:
        return _nicLinkSpeedCache[iface]
    speed_bps = None
    try:
        output = process_util.executeCommandShellWithTimeout('entstat -d %s' % iface, timeout=10)
        for line in output.split('\n'):
            line = line.strip()
            if 'Media Speed Running:' in line:
                parts = line.split(':')
                if len(parts) > 1:
                    tokens = parts[1].strip().split()
                    if tokens:
                        mbps = float(tokens[0])
                        speed_bps = mbps * 1000000.0
                break
    except Exception:
        speed_bps = None
    _nicLinkSpeedCache[iface] = speed_bps
    return speed_bps


_oldNetInfoPerNicMap = None

def GetNetInfo():
    global _oldNetInfoPerNicMap
    timestamp = int(time.time() * 1000)
    iocounters = psutil.net_io_counters(pernic=True)

    newMap = {}
    for deviceId, counter in iocounters.items():
        raw = _NetPerfRaw()
        raw.deviceId = deviceId
        raw.readByteCount = int(counter.bytes_recv)
        raw.readCount = int(counter.packets_recv)
        raw.readErrorCount = int(counter.errin)
        raw.readDroppedCount = int(counter.dropin)
        raw.writeByteCount = int(counter.bytes_sent)
        raw.writeCount = int(counter.packets_sent)
        raw.writeErrorCount = int(counter.errout)
        raw.writeDroppedCount = int(counter.dropout)
        raw.timestamp = timestamp
        newMap[deviceId] = raw

    if not _oldNetInfoPerNicMap:
        _oldNetInfoPerNicMap = newMap
        return None

    total_rx_bps = 0.0
    total_tx_bps = 0.0
    total_rx_pps = 0.0
    total_tx_pps = 0.0
    max_usage = -1.0
    max_nic_name = None

    for deviceId, raw in newMap.items():
        if deviceId not in _oldNetInfoPerNicMap:
            continue
        old = _oldNetInfoPerNicMap[deviceId]
        timediff = raw.timestamp - old.timestamp
        if timediff <= 0:
            continue
        rx_bps = float(raw.readByteCount - old.readByteCount) / timediff * 1000.0 * 8
        tx_bps = float(raw.writeByteCount - old.writeByteCount) / timediff * 1000.0 * 8
        rx_pps = float(raw.readCount - old.readCount) / timediff * 1000.0
        tx_pps = float(raw.writeCount - old.writeCount) / timediff * 1000.0
        total_rx_bps += rx_bps
        total_tx_bps += tx_bps
        total_rx_pps += rx_pps
        total_tx_pps += tx_pps

        if not deviceId.startswith('lo'):
            link_speed = GetNicLinkSpeed(deviceId)
            if link_speed and link_speed > 0:
                usage = (rx_bps + tx_bps) / link_speed * 100.0
                if usage > max_usage:
                    max_usage = usage
                    max_nic_name = deviceId

    _oldNetInfoPerNicMap = newMap

    n = _NetPerf()
    n.trafficIn = total_rx_bps
    n.trafficOut = total_tx_bps
    n.packetIn = total_rx_pps
    n.packetOut = total_tx_pps
    if max_usage >= 0:
        n.linkUsageMax = min(max_usage, 100.0)
        n.nicName = max_nic_name
    return n

def test():
    print 'test'
    print GetNicUtil()
    time.sleep(5)
    for netperf in GetNicUtil():
        print netperf

if __name__ == '__main__':
    test()
