#include "Debug.h"
#include "LuaTools.h"
#include "PluginManager.h"
#include "PluginLua.h"

#include "modules/Buildings.h"
#include "modules/Maps.h"
#include "modules/Units.h"

#include "df/building.h"
#include "df/creature_raw.h"
#include "df/item.h"
#include "df/map_block.h"
#include "df/tile_occupancy.h"
#include "df/unit.h"
#include "df/world.h"

namespace DFHack
{
    DBG_DECLARE(fixoccupancy, log, DebugCategory::LINFO);
}

using namespace DFHack;
using std::string;
using std::vector;

DFHACK_PLUGIN("fix-occupancy");

REQUIRE_GLOBAL(world);

DFhackCExport command_result plugin_init (color_ostream &out, vector<PluginCommand> &commands) {
    return CR_OK;
}

/////////////////////////////////////////////////////
// Lua API
//

struct Expected {
private:
    int32_t dim_x, dim_y, dim_z;
    size_t size;
    df::tile_occupancy * occ_buf;
    df::building ** bld_buf;
    std::unordered_map<df::map_block *, std::set<int32_t>> block_items_buf;

public:
    Expected() {
        Maps::getTileSize(dim_x, dim_y, dim_z);
        size = dim_x * dim_y * dim_z;
        occ_buf = (df::tile_occupancy *)calloc(size, sizeof(*occ_buf));
        bld_buf = (df::building **)calloc(size, sizeof(*bld_buf));
    }

    ~Expected() {
        free(occ_buf);
        free(bld_buf);
    }

    size_t get_size() const {
        return size;
    }

    df::tile_occupancy * occ(int32_t x, int32_t y, int32_t z) {
        size_t off = get_offset(x, y, z);
        if (off < size)
            return &occ_buf[off];
        return nullptr;
    }
    df::tile_occupancy * occ(const df::coord & pos) {
        return occ(pos.x, pos.y, pos.z);
    }

    df::building ** bld(int32_t x, int32_t y, int32_t z) {
        size_t off = get_offset(x, y, z);
        if (off < size)
            return &bld_buf[off];
        return nullptr;
    }
    df::building ** bld(const df::coord & pos) {
        return bld(pos.x, pos.y, pos.z);
    }

    std::set<int32_t> * block_items(df::map_block * block) {
        if (!block)
            return nullptr;
        return &block_items_buf[block];
    }
    std::set<int32_t> * block_items(int32_t x, int32_t y, int32_t z) {
        return block_items(Maps::getTileBlock(x, y, z));
    }
    std::set<int32_t> * block_items(const df::coord & pos) {
        return block_items(pos.x, pos.y, pos.z);
    }

private:
    size_t get_offset(int32_t x, int32_t y, int32_t z) {
        return (dim_x * dim_y * z) + (dim_x * y) + x;
    }
};

static void scan_building(color_ostream &out, df::building * bld, Expected & expected) {
    for (int y = bld->y1; y <= bld->y2; ++y) {
        for (int x = bld->x1; x <= bld->x2; ++x) {
            if (!Buildings::containsTile(bld, df::coord2d(x, y)))
                continue;
            auto expected_bld = expected.bld(x, y, bld->z);
            if (bld->isSettingOccupancy() && expected_bld) {
                if (*expected_bld) {
                    WARN(log,out).print("Buildings overlap at ({}, {}, {}); please manually remove overlapping building."
                        " Run ':lua dfhack.gui.revealInDwarfmodeMap({}, {}, {}, true, true)' to zoom to the tile.\n",
                        x, y, bld->z, x, y, bld->z);
                }
                *expected_bld = bld;
            }
            if (auto expected_occ = expected.occ(x, y, bld->z)) {
                auto bld_occ = df::tile_building_occ::Impassable;
                if (auto block_occ = Maps::getTileOccupancy(x, y, bld->z))
                    bld_occ = block_occ->bits.building;
                expected_occ->bits.building = bld_occ;
            }
        }
    }
}

static void scan_unit(df::unit * unit, Expected & expected) {
    if (unit->flags1.bits.caged || unit->flags1.bits.inactive || unit->flags1.bits.rider)
        return;
    if (auto craw = df::creature_raw::find(unit->race); craw && craw->flags.is_set(df::creature_raw_flags::EQUIPMENT_WAGON)) {
        for (int y = unit->pos.y - 1; y <= unit->pos.y + 1; ++y) {
            for (int x = unit->pos.x - 1; x <= unit->pos.x + 1; ++x) {
                if (auto expected_occ = expected.occ(x, y, unit->pos.z)) {
                    expected_occ->bits.unit = expected_occ->bits.unit || !unit->flags1.bits.on_ground;
                    expected_occ->bits.unit_grounded = expected_occ->bits.unit_grounded || unit->flags1.bits.on_ground;
                }
            }
        }
    } else if (auto expected_occ = expected.occ(unit->pos)) {
        expected_occ->bits.unit = expected_occ->bits.unit || !unit->flags1.bits.on_ground;
        expected_occ->bits.unit_grounded = expected_occ->bits.unit_grounded || unit->flags1.bits.on_ground;
    }
}

static void scan_item(df::item * item, Expected & expected) {
    if (!item->flags.bits.on_ground)
        return;
    auto pos = Items::getPosition(item);
    if (auto block_items = expected.block_items(pos))
        block_items->emplace(item->id);
    if (auto expected_occ = expected.occ(pos))
        expected_occ->bits.item = true;
}

static void normalize_item_vector(color_ostream &out, df::map_block *block, bool dry_run) {
    bool needs_sorting = false;
    int32_t prev_id = -1;
    for (auto item_id : block->items) {
        // do we need to handle duplicates as well?
        if (item_id < prev_id) {
            needs_sorting = true;
            break;
        }
        prev_id = item_id;
    }
    if (needs_sorting) {
        INFO(log,out).print("{} item list for map block at ({}, {}, {})\n",
            dry_run ? "would fix" : "fixing", block->map_pos.x, block->map_pos.y, block->map_pos.z);
        if (!dry_run)
            std::sort(block->items.begin(), block->items.end());
    }
}

static void reconcile_map_tile(color_ostream &out, df::building * bld, const df::tile_occupancy & expected_occ,
    df::tile_occupancy & block_occ, bool dry_run, int x, int y, int z)
{
    // clear building occupancy if there is no building there
    if (expected_occ.bits.building == df::tile_building_occ::None && block_occ.bits.building != df::tile_building_occ::None) {
        INFO(log,out).print("{} building occupancy at ({}, {}, {})\n",
            dry_run ? "would fix" : "fixing", x, y, z);
        if (!dry_run)
            block_occ.bits.building = df::tile_building_occ::None;
    }

    // recalculate bulding occupancy if there *is* a building there
    if (bld) {
        auto prev_occ = block_occ.bits.building;
        bld->updateOccupancy(x, y);
        // if this resets the occupancy to Dynamic, trust the original value
        if (block_occ.bits.building == df::tile_building_occ::Dynamic)
            block_occ.bits.building = prev_occ;
        else if (prev_occ != block_occ.bits.building) {
            INFO(log,out).print("{} building occupancy at ({}, {}, {})\n",
                dry_run ? "would fix" : "fixing", x, y, z);
            if (dry_run)
                block_occ.bits.building = prev_occ;
        }
    }

    // clear unit occupancy if there are no units there
    if (!expected_occ.bits.unit && block_occ.bits.unit) {
        INFO(log,out).print("{} standing unit occupancy at ({}, {}, {})\n",
            dry_run ? "would fix" : "fixing", x, y, z);
        if (!dry_run)
            block_occ.bits.unit = false;
    }
    if (!expected_occ.bits.unit_grounded && block_occ.bits.unit_grounded) {
        INFO(log,out).print("{} grounded unit occupancy at ({}, {}, {})\n",
            dry_run ? "would fix" : "fixing", x, y, z);
        if (!dry_run)
            block_occ.bits.unit_grounded = false;
    }

    // clear item occupancy if there are no items there
    if (!expected_occ.bits.item && block_occ.bits.item) {
        INFO(log,out).print("{} item occupancy at ({}, {}, {})\n",
            dry_run ? "would fix" : "fixing", x, y, z);
        if (!dry_run)
            block_occ.bits.item = false;
    }
}

static void fix_tile(color_ostream &out, df::coord pos, bool dry_run) {
    auto occ = Maps::getTileOccupancy(pos);
    auto block = Maps::getTileBlock(pos);
    if (!occ || !block) {
        WARN(log,out).print("invalid tile: ({}, {}, {})\n", pos.x, pos.y, pos.z);
        return;
    }

    Expected expected;

    // building occupancy (scan all buildings since we can't depend on Buildings::findAtTile)
    size_t num_buildings = 0;
    for (auto bld : world->buildings.all) {
        if (Buildings::containsTile(bld, pos)) {
            ++num_buildings;
            scan_building(out, bld, expected);
        }
    }

    // unit occupancy (make sure we pick up wagons that might be overlapping the tile)
    vector<df::unit *> units;
    Units::getUnitsInBox(units, pos.x-1, pos.y-1, pos.z, pos.x+1, pos.y+1, pos.z);
    size_t num_units = units.size();
    for (auto unit : units)
        scan_unit(unit, expected);

    // check/fix order of map block item vector
    normalize_item_vector(out, block, dry_run);

    // scan for items at that position in the map block
    // this won't catch issues with items that are members of an incorrect map block. use fix_map for that.
    size_t num_items = 0;
    for (auto item_id : block->items) {
        if (auto item = df::item::find(item_id)) {
            ++num_items;
            scan_item(item, expected);
        }
    }

    // check/fix occupancy
    auto expected_occ = expected.occ(pos);
    auto expected_bld = expected.bld(pos);
    if (expected_bld && expected_occ)
        reconcile_map_tile(out, *expected_bld, *expected_occ, block->occupancy[pos.x&15][pos.y&15], dry_run, pos.x, pos.y, pos.z);

    INFO(log,out).print("verified {} building(s), {} unit(s), {} item(s), 1 map block(s), and 1 map tile(s)\n",
        num_buildings, num_units, num_items);
}

static void reconcile_block_items(color_ostream &out, std::set<int32_t> * expected_items, df::map_block * block, bool dry_run) {
    vector<int32_t> & block_items = block->items;

    if (!expected_items) {
        if (block_items.size()) {
            INFO(log,out).print("{} stale item references in map block at ({}, {}, {})\n",
                dry_run ? "would fix" : "fixing", block->map_pos.x, block->map_pos.y, block->map_pos.z);
            if (!dry_run)
                block_items.resize(0);
        }
        return;
    }

    if (!std::equal(expected_items->begin(), expected_items->end(), block_items.begin(), block_items.end())) {
        INFO(log,out).print("{} stale item references in map block at ({}, {}, {})\n",
            dry_run ? "would fix" : "fixing", block->map_pos.x, block->map_pos.y, block->map_pos.z);
        if (!dry_run) {
            block_items.resize(expected_items->size());
            std::copy(expected_items->begin(), expected_items->end(), block_items.begin());
        }
    }
}

static void fix_map(color_ostream &out, bool dry_run) {
    static const uint32_t occ_mask = df::tile_occupancy::mask_building | df::tile_occupancy::mask_unit |
        df::tile_occupancy::mask_unit_grounded | df::tile_occupancy::mask_item;

    Expected expected;

    // set expected building occupancy
    for (auto bld : world->buildings.all)
        scan_building(out, bld, expected);

    // set expected unit occupancy
    for (auto unit : world->units.active)
        scan_unit(unit, expected);

    // set expected item occupancy
    for (auto item : world->items.other.IN_PLAY)
        scan_item(item, expected);

    // check against expected values and fix
    for (auto block : world->map.map_blocks) {
        // check/fix order of map block item vector
        normalize_item_vector(out, block, dry_run);
        reconcile_block_items(out, expected.block_items(block), block, dry_run);

        // reconcile occupancy of each tile against expected values
        int z = block->map_pos.z;
        for (int yoff = 0; yoff < 16; ++yoff) {
            int y = block->map_pos.y + yoff;
            for (int xoff = 0; xoff < 16; ++xoff) {
                int x = block->map_pos.x + xoff;
                auto expected_occ = expected.occ(x, y, z);
                auto expected_bld = expected.bld(x, y, z);
                if (!expected_occ || !expected_bld) {
                    TRACE(log,out).print("pos out of bounds ({}, {}, {})\n", x, y, z);
                    continue;
                }
                df::tile_occupancy &block_occ = block->occupancy[xoff][yoff];
                if (*expected_bld || (expected_occ->whole & occ_mask) != (block_occ.whole & occ_mask)) {
                    DEBUG(log,out).print("reconciling occupancy at ({}, {}, {}) (bld={}, 0x{:x} ?= 0x{:x})\n",
                        x, y, z, static_cast<void*>(*expected_bld), expected_occ->whole & occ_mask, block_occ.whole & occ_mask);
                    reconcile_map_tile(out, *expected_bld, *expected_occ, block_occ, dry_run, x, y, z);
                }
            }
        }
    }

    INFO(log,out).print("verified {} buildings, {} units, {} items, {} map blocks, and {} map tiles\n",
        world->buildings.all.size(), world->units.active.size(), world->items.other.IN_PLAY.size(),
        world->map.map_blocks.size(), expected.get_size());
}

DFHACK_PLUGIN_LUA_FUNCTIONS{
    DFHACK_LUA_FUNCTION(fix_tile),
    DFHACK_LUA_FUNCTION(fix_map),
    DFHACK_LUA_END
};