#include <cstring>
#include <fstream>
#include <filesystem>

#include <WorldComponent/Wiring/Wiring.hpp>
#include <Chunk/Chunk.hpp>
#include <World/LocalWorld/LocalWorld.hpp>

Wiring::GraphWalker::GraphWalker(std::vector<Port> &ports, uint64_t startingPort) :
    ports(ports),
    startingPort(startingPort)
{}

std::optional<uint64_t> Wiring::GraphWalker::findPortIndex(glm::ivec3 pos, Side side)
{
    return findPortIndex(pos, side, ports);
}

std::optional<uint64_t> Wiring::GraphWalker::findPortIndex(glm::ivec3 pos, Side side, std::vector<Port> &ports)
{
    uint64_t i;
    for (i=0; i<ports.size(); i++)
    {
        Port &port = ports[i];
        if (port.pos == pos && port.side == side)
        {
            return i;
        }
    }

    return std::nullopt;
}

void Wiring::GraphWalker::onConnectionFound(uint64_t targetPort, uint32_t distance)
{
    bool foundConnection = false;
    for (Connection &connection : connectionsFound)
    {
        if (connection.targetPort == targetPort)
        {
            foundConnection = true;

            if (distance < connection.distance)
            {
                connection.distance = distance;
            }

            break;
        }
    }

    if (!foundConnection)
    {
        connectionsFound.push_back({targetPort, distance});
    }
}

bool Wiring::GraphWalker::compareConnections()
{
    std::vector<Connection> &currentConnections = ports[startingPort].connections;
    if (currentConnections.size() != connectionsFound.size())
    {
        return true;
    }

    size_t i;
    for (i=0; i<currentConnections.size(); i++)
    {
        Connection &a = currentConnections[i];
        Connection &b = connectionsFound[i];

        if (a.targetPort != b.targetPort || a.distance != b.distance)
        {
            return true;
        }
    }

    return false;
}

bool Wiring::GraphWalker::compareAndUpdateConnections()
{
    if (compareConnections())
    {
        ports[startingPort].connections = connectionsFound;
        return true;
    }

    return false;
}

void Wiring::ChunkWalker::walkRecur(glm::ivec3 pos, Side sourceSide, uint16_t distance)
{
    if ((pos & ~0xF) != chunk.base)
    {
        return;
    }

    glm::ivec3 offset = pos & 0xF;
    size_t index = (offset.z << 8) | (offset.y << 4) | offset.x;

    if (shortestDistanceMap[index] != 0 && shortestDistanceMap[index] <= distance)
    {
        return;
    }

    shortestDistanceMap[index] = distance;

    uint8_t sideBitmap = std::get<1>(chunk.getBlockAndMeta(offset));
    for (const Side &side : Side::all)
    {
        if (side == sourceSide)
        {
            continue;
        }

        std::optional<uint64_t> targetPortIndex = findPortIndex(pos, side);
        if (targetPortIndex.has_value())
        {
            onConnectionFound(*targetPortIndex, distance);
            continue;
        }

        if ((sideBitmap & side.mask()) == 0)
        {
            continue;
        }

        walkRecur(pos + side.out(), side.opposite(), distance + 1);
    }
}

Wiring::ChunkWalker::ChunkWalker(Chunk &chunk, std::vector<Port> &ports, uint64_t startingPort) :
    GraphWalker(ports, startingPort),
    chunk(chunk)
{
    std::memset(shortestDistanceMap, 0, sizeof(shortestDistanceMap));
}

void Wiring::ChunkWalker::run()
{
    Port &port = ports[startingPort];
    walkRecur(port.pos, port.side, 1);

    std::sort(connectionsFound.begin(), connectionsFound.end(), [](const Connection& a, const Connection& b)
    {
        if (a.distance == b.distance)
        {
            return a.targetPort < b.targetPort;
        }

        return a.distance < b.distance;
    });
}

void Wiring::BranchWalker::walkRecur(glm::ivec3 pos, Side side, uint32_t distance)
{
    if (distance > wireDistanceLimit)
    {
        return;
    }

    uint16_t childIndex = node.getChildIndex(pos);
    std::vector<Port> &childPorts = cells[childIndex];

    std::optional<uint64_t> childPortIndex = findPortIndex(pos, side, childPorts);
    if (!childPortIndex.has_value())
    {
        return;
    }

    Port &port = childPorts[*childPortIndex];
    if (port.lastVisitorId == visitorId && port.visitorDistance <= distance)
    {
        return;
    }

    port.lastVisitorId = visitorId;
    port.visitorDistance = distance;

    for (Connection &connection : port.connections)
    {
        Port &targetPort = childPorts[connection.targetPort];
        std::optional<uint64_t> ourTargetPortIndex = findPortIndex(targetPort.pos, targetPort.side);
        
        if (ourTargetPortIndex.has_value())
        {
            onConnectionFound(*ourTargetPortIndex, distance + connection.distance);
        }
        else
        {
            walkRecur(
                targetPort.pos + targetPort.side.out(),
                targetPort.side.opposite(),
                distance + connection.distance
            );
        }
    }
}

Wiring::BranchWalker::BranchWalker(
    WorldGridNode &node,
    CellTable &cells,
    std::vector<Port> &ports,
    uint64_t startingPort,
    uint64_t visitorId
) : GraphWalker(ports, startingPort),
    node(node),
    cells(cells),
    visitorId(visitorId)
{}

void Wiring::BranchWalker::run()
{
    Port &port = ports[startingPort];
    walkRecur(port.pos, port.side, 0);

    std::sort(connectionsFound.begin(), connectionsFound.end(), [](const Connection& a, const Connection& b)
    {
        if (a.distance == b.distance)
        {
            return a.targetPort < b.targetPort;
        }

        return a.distance < b.distance;
    });
}

void Wiring::rerouteLeaf(LocalWorld &world, std::shared_ptr<WorldGridNode> node)
{
    Chunk &chunk = *node->chunk;

    std::shared_ptr<WorldGridNode> parentNode = node->getParent();
    Wiring &parentWiring = parentNode->getComponent<Wiring>(*this);
    std::vector<Port> &ports = parentWiring.cells[node->getParentIndex()];

    bool haveDifferences = false;

    uint64_t i;
    for (i=0; i<ports.size(); i++)
    {
        ChunkWalker walker(chunk, ports, i);
        glm::ivec3 offset = ports[i].pos & 0xF;
        
        if (isWalkableBlock(world, std::get<0>(chunk.getBlockAndMeta(offset))))
        {
            walker.run();
        }

        if (walker.compareAndUpdateConnections())
        {
            haveDifferences = true;
        }
    }

    if (haveDifferences)
    {
        parentNode->getComponent<Wiring>(*this).requestReroute(world, parentNode);
    }
}

void Wiring::rerouteBranch(LocalWorld &world, std::shared_ptr<WorldGridNode> node)
{
    std::shared_ptr<WorldGridNode> parentNode = node->getParent();
    std::vector<Port> &ports =
        parentNode == nullptr ?
        globalPorts :
        parentNode->getComponent<Wiring>(*this).cells[node->getParentIndex()];

    bool haveDifferences = false;

    uint64_t i;
    for (i=0; i<ports.size(); i++)
    {
        BranchWalker walker(*node, cells, ports, i, nextVisitorId++);
        walker.run();

        if (walker.compareAndUpdateConnections())
        {
            haveDifferences = true;
        }
    }

    if (!haveDifferences)
    {
        return;
    }

    if (parentNode != nullptr)
    {
        parentNode->getComponent<Wiring>(*this).requestReroute(world, parentNode);
    }
    else
    {
        onMachinePortsChanged(world, node, ports);
    }
}

Wiring::Wiring(const uint8_t *data, size_t size)
{
    deserialize(data, size, cells, globalPorts);
}

void Wiring::requestReroute(LocalWorld &world, std::shared_ptr<WorldGridNode> node)
{
    needsReroute = true;
    world.markNodeForUpdate(node);
}

void Wiring::serializeState(std::vector<uint8_t> &outBytes)
{
    if (cells.size() == 0)
    {
        // Nothing to serialize if there are no cells.
        return;
    }

    serializeInto(outBytes, cells, globalPorts);
}

void Wiring::onWireBlockPlaced(LocalWorld &world, std::shared_ptr<WorldGridNode> node, glm::ivec3 pos)
{
    if (node->isRoot())
    {
        // Wires don't create ports at the root node.
        requestReroute(world, node);
        return;
    }

    std::shared_ptr<WorldGridNode> parentNode = node->getParent();
    Wiring &parentWiring = parentNode->getComponent<Wiring>(*this);
    std::vector<Port> &ports = parentWiring.cells[node->getParentIndex()];

    int32_t baseMask = ~((UINT32_C(0x10) << (4 * node->getLevel()))-1);
    for (const Side &side : Side::all)
    {
        glm::ivec3 target = pos + side.out();
        if ((target & baseMask) != (pos & baseMask))
        {
            ports.push_back({pos, side, {}});
        }
    }

    parentWiring.onWireBlockPlaced(world, parentNode, pos);
    requestReroute(world, node);
}

void Wiring::onWireBlockRemoved(LocalWorld &world, std::shared_ptr<WorldGridNode> node, glm::ivec3 pos)
{
    if (node->isRoot())
    {
        // Wires don't create ports at the root node.
        requestReroute(world, node);
        return;
    }

    std::shared_ptr<WorldGridNode> parentNode = node->getParent();
    Wiring &parentWiring = parentNode->getComponent<Wiring>(*this);
    std::vector<Port> &ports = parentWiring.cells[node->getParentIndex()];

    int32_t baseMask = ~((UINT32_C(0x10) << (4 * node->getLevel()))-1);

    size_t i;
    for (i=0; i<ports.size(); i++)
    {
        Port &port = ports[i];
        while ((port.pos == pos) && (((pos + port.side.out()) & baseMask) != (pos & baseMask)))
        {
            if (i != ports.size() - 1)
            {
                std::swap(port, ports[ports.size()-1]);
            }

            ports.resize(ports.size()-1);
            if (i == ports.size())
            {
                break;
            }
        }
    }

    parentWiring.onWireBlockRemoved(world, parentNode, pos);
    requestReroute(world, node);
}

void Wiring::onMachineSidePlaced(LocalWorld &world, std::shared_ptr<WorldGridNode> node, glm::ivec3 pos, Side side)
{
    onMachineSidePlaced(world, node, pos, side, pos);
}

void Wiring::onMachineSidePlaced(
    LocalWorld &world,
    std::shared_ptr<WorldGridNode> node,
    glm::ivec3 pos,
    Side side,
    glm::ivec3 machinePos
)
{
    glm::ivec3 portPos = pos + side.out();
    Side portSide = side.opposite();

    while (node != nullptr)
    {
        std::shared_ptr<WorldGridNode> parentNode = node->getParent();
        std::vector<Port> &ports =
            parentNode == nullptr ?
            node->getComponent<Wiring>(*this).globalPorts :
            parentNode->getComponent<Wiring>(*this).cells[node->getParentIndex()];
        
        if (node->containsPos(portPos))
        {
            ports.push_back({portPos, portSide, {}, machinePos});
            node->getComponent<Wiring>(*this).requestReroute(world, node);
        }

        node = parentNode;
    }
}

void Wiring::onMachineSideRemoved(LocalWorld &world, std::shared_ptr<WorldGridNode> node, glm::ivec3 pos, Side side)
{
    glm::ivec3 portPos = pos + side.out();
    Side portSide = side.opposite();

    while (node != nullptr)
    {
        std::shared_ptr<WorldGridNode> parentNode = node->getParent();
        std::vector<Port> &ports =
            parentNode == nullptr ?
            node->getComponent<Wiring>(*this).globalPorts :
            parentNode->getComponent<Wiring>(*this).cells[node->getParentIndex()];
        
        if (node->containsPos(portPos))
        {
            for (size_t i=0; i<ports.size(); i++)
            {
                Port &port = ports[i];
                if (port.pos == portPos && port.side == portSide)
                {
                    if (i != ports.size()-1)
                    {
                        std::swap(port, ports[ports.size()-1]);
                    }

                    ports.resize(ports.size()-1);
                    node->getComponent<Wiring>(*this).requestReroute(world, node);
                    break;
                }
            }
        }

        node = parentNode;
    }
}

bool Wiring::update(LocalWorld &world, std::shared_ptr<WorldGridNode> node)
{
    if (!needsReroute)
    {
        return false;
    }

    needsReroute = false;

    if (node->getLevel() == 0)
    {
        // We never need to save a leaf node because our wiring is stored in the parent.
        rerouteLeaf(world, node);
        return false;
    }
    else
    {
        // If a branch is being rerouted, it means a child changed, and since children save
        // their routing grid in the parent, we must save.
        rerouteBranch(world, node);
        return true;
    }
}

void Wiring::dump(std::shared_ptr<WorldGridNode> node)
{
    std::string dumpDir = getDumpPath(getDumpDirName());
    std::filesystem::create_directories(dumpDir);

    std::string dumpPath = formatString("$+/lvl$+.dot", dumpDir, node->getLevel());

    std::ofstream ofs(dumpPath);
    ofs << "digraph Level" << node->getLevel() << " {" << std::endl;

    std::shared_ptr<WorldGridNode> parentNode = node->getParent();
    std::vector<Port> &ports =
        parentNode == nullptr ?
        globalPorts :
        parentNode->getComponent<Wiring>(*this).cells[node->getParentIndex()];
    
    for (size_t i=0; i<ports.size(); i++)
    {
        Port &port = ports[i];

        ofs << formatString(
            "  P$+ [label=\"($+, $+, $+ : $+)\"];",
            i, 
            port.pos.x, port.pos.y, port.pos.z,
            port.side
        ) << std::endl;

        for (Connection &connection : port.connections)
        {
            ofs << formatString("  P$+ -> P$+ [label=\"$+\"];", i, connection.targetPort, connection.distance) << std::endl;
        }
    }

    for (auto &[childIndex, childPorts] : cells)
    {
        for (size_t i=0; i<childPorts.size(); i++)
        {
            Port &childPort = childPorts[i];
            std::string ourNodeName = formatString(
                "C_X$+_Y$+_Z$+_S$+",
                (uint32_t) childPort.pos.x,
                (uint32_t) childPort.pos.y,
                (uint32_t) childPort.pos.z,
                childPort.side
            );

            ofs << formatString(
                "  $+ [label=\"($+, $+, $+ : $+)\",style=dashed];",
                ourNodeName, 
                childPort.pos.x, childPort.pos.y, childPort.pos.z,
                childPort.side
            ) << std::endl;

            size_t ourPortIndex;
            for (ourPortIndex=0; ourPortIndex<ports.size(); ourPortIndex++)
            {
                Port &ourPort = ports[ourPortIndex];
                if (ourPort.pos == childPort.pos && ourPort.side == childPort.side)
                {
                    ofs << formatString("  $+ -> P$+ [style=dotted];", ourNodeName, ourPortIndex) << std::endl;
                    ofs << formatString("  P$+ -> $+ [style=dotted];", ourPortIndex, ourNodeName) << std::endl;
                    break;
                }
            }

            if (ourPortIndex == ports.size())
            {
                glm::ivec3 oppositePos = childPort.pos + childPort.side.out();
                std::string oppositeNodeName = formatString(
                    "C_X$+_Y$+_Z$+_S$+",
                    (uint32_t) oppositePos.x,
                    (uint32_t) oppositePos.y,
                    (uint32_t) oppositePos.z,
                    childPort.side.opposite()
                );

                ofs << "  " << ourNodeName << " -> " << oppositeNodeName << " [style=dotted];" << std::endl;
            }

            for (Connection &connection : childPort.connections)
            {
                Port &targetPort = childPorts[connection.targetPort];
                std::string targetNodeName = formatString(
                    "C_X$+_Y$+_Z$+_S$+",
                    (uint32_t) targetPort.pos.x,
                    (uint32_t) targetPort.pos.y,
                    (uint32_t) targetPort.pos.z,
                    targetPort.side
                );

                ofs << formatString(
                    "  $+ -> $+ [label=\"$+\",style=dashed];",
                    ourNodeName,
                    targetNodeName,
                    connection.distance
                ) << std::endl;
            }
        }
    }

    ofs << "}" << std::endl;
    ofs.close();
}