#include <vector>
#include <string>
#include <sstream>
#include <list>

#include "ap-address-manager.h"
#include "ns3/node.h"
#include "ns3/config.h"
#include "ns3/node-list.h"
#include "ns3/ipv4.h"
#include "ns3/simulator.h"
#include "ns3/log.h"

namespace ns3 {
    NS_LOG_COMPONENT_DEFINE("ApAddressManager");
    
    void ApAddressManager::AddAccessPoint(Ptr<const NetDevice> dev, const string& addr_base, unsigned int first_valid_addr) {
        access_point ap;
        ap.address = convert(dev->GetAddress());
        ap.addressBase = addr_base;
        for (unsigned int i = first_valid_addr; i < 255; i++) {
            ap.addresses.push(i);
        }
        ap.ip = dev->GetNode()->GetObject<Ipv4 > ()->GetAddress(2);
        ap.node = dev->GetNode();
        accessPoints[ap.address] = ap;
    }

    void ApAddressManager::AddStations(NetDeviceContainer& devs) {
        for (unsigned int i = 0; i < devs.GetN(); i++) {
            Ptr<Node> n = devs.Get(i)->GetNode();
            unsigned int dev_id = getDeviceId(n, devs.Get(i)->GetTypeId());
            ostringstream path1, path2;
            Ptr<Ipv4> ip = n->GetObject<Ipv4 > ();
            path1 << "/NodeList/" << n->GetId() << "/DeviceList/" << dev_id << "/$ns3::WifiNetDevice/Mac/$ns3::NqstaWifiMac/Assoc";
            path2 << "/NodeList/" << n->GetId() << "/DeviceList/" << dev_id << "/$ns3::WifiNetDevice/Mac/$ns3::NqstaWifiMac/DeAssoc";
            Config::Connect(path1.str(), MakeCallback(&ApAddressManager::assoc, this));
            Config::Connect(path2.str(), MakeCallback(&ApAddressManager::deassoc, this));
        }
    }

    void ApAddressManager::assoc(string context, Mac48Address ap) {
        NS_LOG_FUNCTION_NOARGS();
        map<Mac48Address, access_point>::iterator it = accessPoints.find(ap);
        if (it == accessPoints.end()) return;
        //assign ip
        unsigned int n_addr, node_id;
        unsigned int if_id;
        Ipv4Address old_addr, addr = getAddr(it->first, &n_addr);
        Ipv4Mask old_mask;
        node_id = getNodeId(context);
        Ptr<Node> n = NodeList::GetNode(node_id);
        Ptr<Ipv4> ip = n->GetObject<Ipv4 > ();
        if_id = getDeviceId(n, WifiNetDevice::GetTypeId()) + 1;
        //if_id=1;
        old_addr = ip->GetAddress(if_id);
        old_mask = ip->GetNetworkMask(if_id);
        ip->SetAddress(if_id, addr);
        ip->SetNetworkMask(if_id, Ipv4Mask("255.255.255.0"));
        while (ip->GetNRoutes() > 1) ip->RemoveRoute(1); //delete the old routes
        ip->SetDefaultRoute(it->second.ip, if_id);
        //store the record
        station s;
        s.nodeId = node_id;
        s.accessPoint = ap;
        s.n_addr = n_addr;
        s.addr = addr;
        s.old_addr = old_addr;
        s.old_mask = old_mask;
        registeredStations[node_id] = s;
        //setup multicast
        ip = it->second.node->GetObject<Ipv4 > ();
        for (list<MulticastRoute>::iterator it = multicastRoutes.begin(); it != multicastRoutes.end(); it++) {
            vector<uint32_t> o;
            o.push_back(it->out_interface);
            ip->AddMulticastRoute(it->outgoing ? Ipv4Address(addr) : Ipv4Address::GetAny(), it->group, it->in_interface, o);
        }
        ip = NodeList::GetNode(node_id)->GetObject<Ipv4 > ();
        NS_LOG_DEBUG("at " << Simulator::Now().GetSeconds() << ": assigned ip " << addr << " to station #" << node_id <<
                "whose ip is now " << ip->GetAddress(if_id));
    }

    void ApAddressManager::deassoc(string context, Mac48Address ap) {
        NS_LOG_FUNCTION("de-assoc from " << ap);
        map<Mac48Address, access_point>::iterator it = accessPoints.find(ap);
        if (it == accessPoints.end()) return;
        //assign old ip, free "our" one
        unsigned int node_id, if_id;
        Ipv4Address addr, new_addr;
        node_id = getNodeId(context);
        Ptr<Node> n = NodeList::GetNode(node_id);
        if_id = getDeviceId(n, WifiNetDevice::GetTypeId()) + 1;
        //if_id=1;
        Ptr<Ipv4> ip = n->GetObject<Ipv4 > ();
        map<unsigned int, station>::iterator s_it = registeredStations.find(node_id);
        if (s_it != registeredStations.end()) {
            addr = s_it->second.addr;
            //it->second.addresses.push(s_it->second.n_addr);
            ip->SetAddress(if_id, s_it->second.old_addr);
            new_addr = s_it->second.old_addr;
            ip->SetNetworkMask(if_id, s_it->second.old_mask);
        }
        //remove the record
        registeredStations.erase(s_it);
        //remove multicast routes
        ip = it->second.node->GetObject<Ipv4 > ();
        for (list<MulticastRoute>::iterator lit = multicastRoutes.begin(); lit != multicastRoutes.end(); lit++) {
            ip->RemoveMulticastRoute(lit->outgoing ? addr : Ipv4Address::GetAny(), lit->group, lit->in_interface);
        }
        NS_LOG_DEBUG("at " << Simulator::Now().GetSeconds() << ": removed ip " << addr << " from station #" << node_id <<
                "whose ip is now " << ip->GetAddress(if_id) << endl
                << "The first available n is now " << it->second.addresses.front());
    }

    Ipv4Address ApAddressManager::getAddr(Mac48Address ap, unsigned int *id) {
        map<Mac48Address, access_point>::iterator it = accessPoints.find(ap);
        ostringstream ip;
        ip << it->second.addressBase << it->second.addresses.front();
        if (id != 0) *id = it->second.addresses.front();
        it->second.addresses.pop();
        return Ipv4Address(ip.str().c_str());
    }

    Mac48Address ApAddressManager::convert(const Address& addr) const {
        uint8_t *buf = new uint8_t[addr.GetLength()];
        addr.CopyTo(buf);
        Mac48Address ret;
        ret.CopyFrom(buf);
        delete[] buf;
        return ret;
    }

    unsigned int ApAddressManager::getDeviceId(Ptr<const Node> n, TypeId type) const {
        unsigned int dev_id = 0;
        while (dev_id < n->GetNDevices() && n->GetDevice(dev_id)->GetInstanceTypeId() != type) dev_id++;
        if (dev_id == n->GetNDevices()) {
            NS_LOG_DEBUG("ApAddressManager::getDeviceId failed, returning default 0");
            return 0;
        } else return dev_id;
    }

    unsigned int ApAddressManager::getNodeId(const string& path) const {
        //the path is in the form /NodeList/node-id/...
        //node-id is between end of "/NodeList/" and the next occurrance of "/"
        string no_nid = path.substr(10);
        int pos = no_nid.find('/');
        istringstream s(no_nid.substr(0, pos));
        unsigned int ret;
        s >> ret;
        return ret;
    }

    list<ApAddressManager::MulticastRoute>& ApAddressManager::MulticastRoutes(void) {
        return multicastRoutes;
    }

    ApAddressManager::MulticastRoute::MulticastRoute(bool outgoing, Ipv4Address group, unsigned int in_interface, unsigned int out_interface) {
        this->outgoing = outgoing;
        this->group = group;
        this->in_interface = in_interface;
        this->out_interface = out_interface;
    }
}//ns3