#include "../gen/Def.h"
#include "Gen.h"
#include "Global.h"
#include "Instance.h"
#include "Log.h"
#include "Puzzle.h"

#include "bx/bx.h"
#include "bx/string.h"
#include "imgui.h"
#include <cassert>

namespace
{
    using namespace Gen;
    static constexpr PuzPos Dirs[4]{
        {-1, 0},
        {0, -1},
        {0, 1},
        {1, 0},
    };

    StaticPuzzleData StaticData;
    StaticPuzzleCard InvalidCard;
} // namespace

namespace Puzzle
{
    constexpr float UIPuzBoxSize = 26;
    using namespace Gen;

    void Setup()
    {
        LOG("Setting up static puzzle data");
        LoadStaticPuzzleData();
    }

    StaticPuzzleData& GetStaticPuzzleData()
    {
        return StaticData;
    }

    void LoadStaticPuzzleData()
    {
        Deserializer ser;
        if (ser.Init("game/data/static/puzzle.dat", "SPUZ") && ser.ReadT(StaticData))
        {
            LOG("Successfully loaded static puzzle data!");
        }
        ser.Finish();
    }

    void SaveStaticPuzzleData()
    {
        auto& data = GetStaticPuzzleData();
        Serializer ser;
        if (ser.Init("game/data/static/puzzle.dat", "SPUZ") && ser.WriteT(GetStaticPuzzleData()))
        {
            LOG("Successfully saved static puzzle data!");
        }
        ser.Finish();
    }
    const StaticPuzzleCard& GetCard(StaticPuzzleCardHandle H)
    {
        if (H.Idx == UINT16_MAX)
        {
            LOG_ERROR("Invalid static card handle!");
            return InvalidCard;
        }
        return GetStaticPuzzleData().Cards[H.Idx];
    }

    bool IsValid(StaticPuzzleCardHandle h)
    {
        return h.Idx != UINT16_MAX;
    }

    const char* GetShortNodeName(const PuzzleElementType::Enum node)
    {
        return PuzzleElementType::ShortName[node];
    }

    uint8_t GetRemainingCount(const PuzzleCardStack& stack)
    {
        assert(stack.MaxAvailableCount >= stack.UsedCount);
        return stack.MaxAvailableCount - stack.UsedCount;
    }

    PuzzleElementType::Enum GetNodeAt(const PuzzleData& puz, PuzPos pos)
    {
        assert(pos.X < Puzzle::Config::MaxPuzzleSizeTiles && pos.Y < Puzzle::Config::MaxPuzzleSizeTiles && pos.X >= 0 &&
               pos.Y >= 0);
        int32_t cardIdxX = pos.X / Puzzle::Config::CardSize;
        int32_t cardIdxY = pos.Y / Puzzle::Config::CardSize;
        auto& card = puz.PlacedCards[cardIdxY * Puzzle::Config::MaxPuzzleSizeCards + cardIdxX];
        int32_t offsetX = pos.X - (cardIdxX * Config::CardSize);
        int32_t offsetY = pos.Y - (cardIdxY * Config::CardSize);
        if (offsetX >= 0 && offsetX < Puzzle::Config::CardSize && offsetY >= 0 && offsetY < Puzzle::Config::CardSize &&
            IsValid(card.RefCard))
        {
            PuzzleElementType::Enum cardVal = GetCardNodeAt(GetCard(card.RefCard), card.Rotation, offsetX, offsetY);
            if (cardVal != PuzzleElementType::None) return cardVal;
        }
        return puz.BackgroundTiles[pos.Y * Puzzle::Config::MaxPuzzleSizeTiles + pos.X];
    }

    PuzzleElementType::Enum GetCardNodeAt(const StaticPuzzleCard& card, uint8_t rotation, int8_t x, int8_t y)
    {
        assert(x >= 0 && x < Puzzle::Config::CardSize);
        assert(y >= 0 && y < Puzzle::Config::CardSize);

        uint8_t arrX = x;
        uint8_t arrY = y;
        if (rotation == 1)
        {
            arrX = y;
            arrY = Config::CardSize - x - 1;
        }
        else if (rotation == 2)
        {
            arrX = Config::CardSize - x - 1;
            arrY = Config::CardSize - y - 1;
        }
        else if (rotation == 3)
        {
            arrX = Config::CardSize - y - 1;
            arrY = x;
        }

        return card.Elements[arrY * Puzzle::Config::CardSize + arrX];
    }

    PuzzleElementType::Enum& EditCardNodeAt(StaticPuzzleCard& card, uint8_t rotation, int8_t x, int8_t y)
    {
        assert(x >= 0 && x < Puzzle::Config::CardSize);
        assert(y >= 0 && y < Puzzle::Config::CardSize);

        uint8_t arrX = x;
        uint8_t arrY = y;
        if (rotation == 1)
        {
            arrX = y;
            arrY = Config::CardSize - x - 1;
        }
        else if (rotation == 2)
        {
            arrX = Config::CardSize - x - 1;
            arrY = Config::CardSize - y - 1;
        }
        else if (rotation == 3)
        {
            arrX = Config::CardSize - y - 1;
            arrY = x;
        }

        return card.Elements[arrY * Puzzle::Config::CardSize + arrX];
    }

    bool PuzzleSolver::IsPuzzleSolved(const PuzzleData& puzzle)
    {
        bool IsSolved = true;
        for (uint32_t i = 0; i < puzzle.GoalPositionCount; ++i)
        {
            if (!IsExitSatisfied(puzzle, puzzle.GoalPositions[i]))
            {
                IsSolved = false;
                break;
            }
        }
        return IsSolved;
    }

    bool PuzzleSolver::IsExitSatisfied(const PuzzleData& puzzle, PuzPos pos)
    {
        PuzzleElementType::Enum goalType = GetNodeAt(puzzle, pos);

        uint32_t currentPositionQueueIdx = 0;
        uint32_t positionQueueCount = 0;
        PuzPos positionQueue[Puzzle::Config::MaxTilesInPuzzle];

        positionQueue[0] = pos;
        positionQueueCount++;
        while (currentPositionQueueIdx < positionQueueCount)
        {
            assert(currentPositionQueueIdx < Puzzle::Config::MaxTilesInPuzzle);
            PuzPos currentPos = positionQueue[currentPositionQueueIdx];
            for (PuzPos dir : Dirs)
            {
                PuzPos nextPos = currentPos + dir;
                if (!(nextPos.X >= 0 && nextPos.Y >= 0 && nextPos.X < puzzle.WidthTiles &&
                      nextPos.Y < puzzle.HeightTiles))
                    continue;
                if (IsValidGoalConnection(puzzle, nextPos, currentPos, goalType))
                {
                    if (IsValidSource(GetNodeAt(puzzle, nextPos), goalType))
                    {
                        return true;
                    }
                    bool found = false;
                    for (int32_t i = 0; i < positionQueueCount; ++i)
                    {
                        if (positionQueue[i].X == nextPos.X && positionQueue[i].Y == nextPos.Y)
                        {
                            found = true;
                            break;
                        }
                    }
                    if (!found)
                    {
                        positionQueue[positionQueueCount] = nextPos;
                        positionQueueCount++;
                    }
                }
            }
            currentPositionQueueIdx++;
        }
        return false;
    }

    bool PuzzleSolver::IsValidGoalConnection(const PuzzleData& puzzle,
                                             PuzPos flowFrom,
                                             PuzPos flowTo,
                                             PuzzleElementType::Enum goalType)
    {
        auto from = GetNodeAt(puzzle, flowFrom);
        auto to = GetNodeAt(puzzle, flowTo);

        if (goalType == PuzzleElementType::WaterGoal)
        {
            return from == PuzzleElementType::WaterIn || from == PuzzleElementType::WaterChannel ||
                   from == PuzzleElementType::WaterGoal || from == PuzzleElementType::Bridge;
        }
        else if (goalType == PuzzleElementType::ElectricGoal)
        {
            return from == PuzzleElementType::ElectricIn || from == PuzzleElementType::ElectricGoal ||
                   from == PuzzleElementType::None;
        }
        assert(false);
        return false;
    }
    bool PuzzleSolver::IsValidSource(PuzzleElementType::Enum sourceType, PuzzleElementType::Enum goalType)
    {
        return (sourceType == PuzzleElementType::WaterIn && goalType == PuzzleElementType::WaterGoal) ||
               (sourceType == PuzzleElementType::ElectricIn && goalType == PuzzleElementType::ElectricGoal);
    }
    void DrawCard(const StaticPuzzleCard& card, uint8_t rotation, ImVec2 pos)
    {
        auto& drawList = *ImGui::GetWindowDrawList();
        for (int32_t y = 0; y < Puzzle::Config::CardSize; ++y)
        {
            for (int32_t x = 0; x < Puzzle::Config::CardSize; ++x)
            {
                ImGui::SetCursorScreenPos({pos.x + x * UIPuzBoxSize, pos.y + y * UIPuzBoxSize});
                ImGui::Text("%s", GetShortNodeName(GetCardNodeAt(card, rotation, x, y)));
            }
        }
        for (int32_t y = 0; y <= Puzzle::Config::CardSize; ++y)
        {
            ImVec2 linePos = {pos.x, pos.y + y * UIPuzBoxSize};
            drawList.AddLine(linePos,
                             {linePos.x + Puzzle::Config::CardSize * UIPuzBoxSize, linePos.y},
                             y % Puzzle::Config::CardSize == 0 ? 0xFFFFFFFF : 0x11FFFFFF);
        }
        for (int32_t x = 0; x <= Puzzle::Config::CardSize; ++x)
        {
            ImVec2 linePos = {pos.x + x * UIPuzBoxSize, pos.y};
            drawList.AddLine(linePos,
                             {linePos.x, linePos.y + Puzzle::Config::CardSize * UIPuzBoxSize},
                             x % Puzzle::Config::CardSize == 0 ? 0xFFFFFFFF : 0x11FFFFFF);
        }
        ImGui::SetCursorScreenPos(pos);
        ImGui::InvisibleButton("cardbn",
                               {Puzzle::Config::CardSize * UIPuzBoxSize, Puzzle::Config::CardSize * UIPuzBoxSize});
    }
    void RotateCard(PlacedPuzzleCard& card)
    {
        card.Rotation += 1;
        if (card.Rotation >= 4) card.Rotation = 0;
    }

    void WritePuzzleFilePath(char* buf, int32_t bufSize, uint16_t puzID)
    {
        bx::snprintf(buf, bufSize, "%s/%u.pzl", Puzzle::PuzzleFileDir, puzID);
    }

    bool ReturnPlacedCard(PuzzleData& obj, PuzPos targetPos)
    {
        auto& placedCard = obj.PlacedCards[targetPos.Y * Config::MaxPuzzleSizeCards + targetPos.X];
        if (IsValid(placedCard.RefCard))
        {
            if (placedCard.IsLocked)
            {
                LOG_WARN("Card at %i %i is locked!", targetPos.X, targetPos.Y);
                return false;
            }
            bool found = false;
            for (int32_t i = 0; i < obj.AvailableCardCount; ++i)
            {
                if (obj.AvailableCards[i].RefCard.Idx == placedCard.RefCard.Idx && obj.AvailableCards[i].UsedCount > 0)
                {
                    obj.AvailableCards[i].UsedCount--;
                    found = true;
                    break;
                }
            }
            if (!found)
            {
                LOG_ERROR("Failed to find available card to return placed card of type %u to!", placedCard.RefCard.Idx);
                return false;
            }
            placedCard.RefCard = {};
        }
        return true;
    }

    bool DragAvailableCardTo(PuzzleData& obj, PuzPos targetPos, int32_t availIdx, uint8_t rotation)
    {
        if (availIdx >= obj.AvailableCardCount)
        {
            LOG_ERROR("Invalid drag with avail idx %i!", availIdx);
            return false;
        }
        if (!ReturnPlacedCard(obj, targetPos)) return false;

        auto& draggedCard = obj.AvailableCards[availIdx];
        draggedCard.UsedCount++;

        auto& placedCard = obj.PlacedCards[targetPos.Y * Config::MaxPuzzleSizeCards + targetPos.X];
        placedCard.RefCard = draggedCard.RefCard;
        placedCard.IsLocked = false;
        placedCard.Position = targetPos;
        placedCard.Rotation = rotation;
        return true;
    }

    bool RenderDebugUI(PuzzleData& obj)
    {
        bool dataChanged = false;
        uint8_t debugRot = Game::GetInstance().DebugData.DebugCardRotation;

        bool isVisible = true;
        if (ImGui::Begin("Puzzle", &isVisible))
        {
            if (ImGui::Button("Delete Puzzle"))
            {
                char filepath[128]{0};
                WritePuzzleFilePath(filepath, sizeof(filepath), obj.ID);
                bx::remove(filepath);
                obj.ID = UINT16_MAX;
                ImGui::End();
                return false;
            }
            ImGui::SameLine();
            if (ImGui::Button("Reload"))
            {
                // TODO
            }
            int32_t val = obj.ID;
            if (ImGui::InputInt("ID", &val))
            {
                obj.ID = val;
                dataChanged = true;
            }
            dataChanged = ImGui::InputText("Name", obj.PuzzleName, sizeof(obj.PuzzleName)) || dataChanged;

            int32_t W = obj.WidthTiles;
            int32_t H = obj.HeightTiles;
            ImGui::PushID("width");
            ImGui::SetNextItemWidth(40);
            if (ImGui::DragInt("", &W, 0.3f, 0, Puzzle::Config::MaxPuzzleSizeTiles))
            {
                obj.WidthTiles = uint8_t(W);
                dataChanged = true;
            }
            ImGui::PopID();
            ImGui::SameLine();
            ImGui::Text("x");
            ImGui::SameLine();
            ImGui::SetNextItemWidth(40);
            ImGui::PushID("height");
            if (ImGui::DragInt("", &H, 0.3f, 0, Puzzle::Config::MaxPuzzleSizeTiles))
            {
                obj.HeightTiles = uint8_t(H);
                dataChanged = true;
            }
            ImGui::PopID();

            ImVec2 puzCursorStart = ImGui::GetCursorScreenPos();
            auto& drawList = *ImGui::GetWindowDrawList();

            obj.GoalPositionCount = 0;
            for (int32_t y = 0; y < obj.HeightTiles; ++y)
            {
                ImGui::PushID(y);
                for (int32_t x = 0; x < obj.WidthTiles; ++x)
                {
                    ImGui::PushID(x);
                    PuzPos nodePos = {int8_t(x), int8_t(y)};
                    auto node = GetNodeAt(obj, nodePos);
                    if (node == PuzzleElementType::WaterGoal)
                    {
                        obj.GoalPositions[obj.GoalPositionCount] = nodePos;
                        obj.GoalPositionCount++;
                    }

                    ImVec2 pos = ImVec2{puzCursorStart.x + x * UIPuzBoxSize + 5, puzCursorStart.y + y * UIPuzBoxSize};
                    ImGui::SetCursorScreenPos(pos);
                    ImGui::Text("%s", GetShortNodeName(node));
                    ImGui::SetCursorScreenPos(pos);
                    if (x % Puzzle::Config::CardSize == 0 && y % Puzzle::Config::CardSize == 0)
                    {
                        int32_t cardX = x / Config::CardSize;
                        int32_t cardY = y / Config::CardSize;
                        PuzPos cardPos = {int8_t(cardX), int8_t(cardY)};
                        auto& placedCard = obj.PlacedCards[cardY * Config::MaxPuzzleSizeCards + cardX];
                        ImGui::InvisibleButton("bn", {UIPuzBoxSize * 2, UIPuzBoxSize * 2});
                        if (ImGui::IsItemClicked(ImGuiMouseButton_Right))
                        {
                            placedCard.Rotation += 1;
                            if (placedCard.Rotation >= 4) placedCard.Rotation = 0;
                            dataChanged = true;
                        }
                        if (ImGui::IsItemClicked(ImGuiMouseButton_Middle) && IsValid(placedCard.RefCard))
                        {
                            if (placedCard.IsLocked)
                            {
                                placedCard.IsLocked = false;
                                placedCard.RefCard = {};
                                dataChanged = true;
                            }
                            else
                            {
                                if (!ReturnPlacedCard(obj, cardPos))
                                {
                                    placedCard.IsLocked = false;
                                    placedCard.RefCard = {};
                                    dataChanged = true;
                                }
                            }
                        }
                        if (!placedCard.IsLocked && IsValid(placedCard.RefCard))
                        {
                            ImVec2 s = {puzCursorStart.x + x * UIPuzBoxSize, puzCursorStart.y + y * UIPuzBoxSize};
                            drawList.AddRectFilled(s,
                                                   {s.x + UIPuzBoxSize * Puzzle::Config::CardSize,
                                                    s.y + UIPuzBoxSize * Puzzle::Config::CardSize},
                                                   0x33FFFFFF);
                        }
                        if (ImGui::BeginDragDropTarget())
                        {
                            if (const ImGuiPayload* payload = ImGui::AcceptDragDropPayload("cardtype", 0))
                            {
                                uint32_t CardIdx = *reinterpret_cast<uint32_t*>(payload->Data);
                                placedCard = {};
                                placedCard.RefCard = {(uint16_t)CardIdx};
                                placedCard.Rotation = debugRot;
                                placedCard.Position = cardPos;
                                placedCard.IsLocked = true;
                                dataChanged = true;
                            }
                            if (const ImGuiPayload* payload = ImGui::AcceptDragDropPayload("availcard", 0))
                            {
                                DragAvailableCardTo(obj,
                                                    cardPos,
                                                    *(int32_t*)payload->Data,
                                                    Game::GetInstance().DebugData.DebugCardRotation);
                            }
                            ImGui::EndDragDropTarget();
                        }
                    }
                    ImGui::PopID();
                }
                ImGui::PopID();
            }

            for (int32_t y = 0; y <= obj.HeightTiles; ++y)
            {
                ImVec2 linePos = {puzCursorStart.x, puzCursorStart.y + y * UIPuzBoxSize};
                drawList.AddLine(linePos,
                                 {linePos.x + obj.WidthTiles * UIPuzBoxSize, linePos.y},
                                 y % Puzzle::Config::CardSize == 0 ? 0xFFFFFFFF : 0x11FFFFFF);
            }
            for (int32_t x = 0; x <= obj.WidthTiles; ++x)
            {
                ImVec2 linePos = {puzCursorStart.x + x * UIPuzBoxSize, puzCursorStart.y};
                drawList.AddLine(linePos,
                                 {linePos.x, linePos.y + obj.HeightTiles * UIPuzBoxSize},
                                 x % Puzzle::Config::CardSize == 0 ? 0xFFFFFFFF : 0x11FFFFFF);
            }

            ImVec2 pos = ImVec2{puzCursorStart.x + obj.WidthTiles * UIPuzBoxSize + 5,
                                puzCursorStart.y + obj.HeightTiles * UIPuzBoxSize};

            ImGui::SetCursorScreenPos(pos);
            ImGui::NewLine();
            ImGui::Separator();

            PuzzleSolver solver;
            ImGui::Text("Solved: %s", solver.IsPuzzleSolved(obj) ? "yes!" : "no.");

            bool bAvailOpen =
                ImGui::TreeNodeEx("Available Cards", ImGuiTreeNodeFlags_DefaultOpen | ImGuiTreeNodeFlags_Framed);
            if (ImGui::BeginDragDropTarget())
            {
                if (const ImGuiPayload* payload = ImGui::AcceptDragDropPayload("cardtype", 0))
                {
                    uint32_t CardIdx = *reinterpret_cast<uint32_t*>(payload->Data);
                    bool found = false;
                    for (int32_t i = 0; i < obj.AvailableCardCount; ++i)
                    {
                        if (obj.AvailableCards[i].RefCard.Idx == CardIdx)
                        {
                            obj.AvailableCards[i].MaxAvailableCount++;
                            found = true;
                            break;
                        }
                    }
                    if (!found)
                    {
                        if (obj.AvailableCardCount == BX_COUNTOF(obj.AvailableCards))
                        {
                            LOG_ERROR("Available card limit reached!");
                        }
                        else
                        {
                            auto& Card = obj.AvailableCards[obj.AvailableCardCount];
                            Card.RefCard.Idx = CardIdx;
                            Card.MaxAvailableCount = 1;
                            obj.AvailableCardCount++;
                        }
                    }
                    dataChanged = true;
                }
                ImGui::EndDragDropTarget();
            }
            if (bAvailOpen)
            {
                ImVec2 startPos = ImGui::GetCursorScreenPos();
                for (uint32_t i = 0; i < obj.AvailableCardCount; ++i)
                {
                    ImVec2 localPos = {startPos.x + i * 60.0f, startPos.y};
                    ImGui::SetCursorScreenPos(localPos);
                    ImGui::PushID(i);
                    auto& card = obj.AvailableCards[i];
                    DrawCard(GetCard(card.RefCard), debugRot, ImGui::GetCursorScreenPos());
                    int displayCount = card.MaxAvailableCount - card.UsedCount;
                    if (displayCount > 0 && ImGui::BeginDragDropSource())
                    {
                        ImGui::SetDragDropPayload("availcard", &i, sizeof(i));
                        DrawCard(GetCard(card.RefCard), debugRot, ImGui::GetCursorScreenPos());
                        ImGui::EndDragDropSource();
                    }

                    ImGui::SetCursorScreenPos({localPos.x, localPos.y + 55.0f});
                    ImGui::SetNextItemWidth(35);
                    ImGui::PushID("carc");
                    if (ImGui::DragInt("", &displayCount, 0.25f))
                    {
                        int diff = displayCount - (card.MaxAvailableCount - card.UsedCount);
                        card.MaxAvailableCount = bx::max(0, card.MaxAvailableCount + diff);
                        dataChanged = true;
                    }
                    ImGui::PopID();
                    ImGui::SameLine(0, 3);
                    if (ImGui::Button("x"))
                    {
                        if (i < obj.AvailableCardCount - 1)
                        {
                            obj.AvailableCards[i] = obj.AvailableCards[obj.AvailableCardCount - 1];
                        }
                        obj.AvailableCardCount--;
                    }
                    ImGui::PopID();
                }
                ImGui::TreePop();
            }
        }
        ImGui::End();

        if (dataChanged)
        {
            char path[128]{0};
            WritePuzzleFilePath(path, sizeof(path), obj.ID);
            Serializer ser;
            if (ser.Init(path, "PZZL") && ser.WriteT(obj))
            {
                LOG("Saved to %s", path);
            }
            else
            {
                LOG_ERROR("Failed to save to %s", path);
            }
            ser.Finish();
        }

        return isVisible;
    }
} // namespace Puzzle