#include "TechTreeImproved.h"

using json = nlohmann::json;

TechTreeImproved::TechTreeImproved() { }

sc2::UNIT_TYPEID string_to_id(const std::string & str)
{
    return static_cast<sc2::UNIT_TYPEID>(std::stoi(str));
}

std::string id_to_string(int id)
{
    return sc2::UnitTypeToName(static_cast<sc2::UNIT_TYPEID>(id));
}

void add_requirement(json & requirement, BuildAlternative & alternative)
{
    std::string type = requirement["type"];
    if (type == "and")
    {
        for (auto & subrequirement : requirement["operands"])
        {
            add_requirement(subrequirement, alternative);
        }
    }
    else if (type == "unitCount")
    {
        if (requirement["state"] == "CompleteOnlyAtUnit")
        {
            for (auto & unit : requirement["unit"])
            {
                //std::cout << "Addon: " << sc2::UnitTypeToName(static_cast<sc2::UNIT_TYPEID>(unit)) << " (" << unit << ")" << std::endl;
                alternative.addons_needed.push_back(static_cast<sc2::UNIT_TYPEID>(unit));
            }
        }
        else if (requirement["state"] == "CompleteOnly")
        {
            for (auto & unit : requirement["unit"])
            {
                //std::cout << "Just building: " << sc2::UnitTypeToName(static_cast<sc2::UNIT_TYPEID>(unit)) << " (" << unit << ")" << std::endl;
                alternative.buildings_needed.push_back(static_cast<sc2::UNIT_TYPEID>(unit));
            }
        }
        else
        {
            //std::cout << "Unsupported" << std::endl;
            //std::cout << requirement << std::endl;
        }
    }
}

BuildAlternative parse_build_alternative(json & build_item)
{
    BuildAlternative alternative{ build_item["unit"] };

    if (build_item.find("requires") != build_item.end())
    {
        //std::cout << "Building unit: " << id_to_string(build_item["unit"]) << " requires" << std::endl;
        auto & requires = build_item["requires"][0];
        add_requirement(requires, alternative);
    }
    return alternative;
}

std::pair<sc2::UNIT_TYPEID, std::vector<BuildAlternative>> parse_unit(json::iterator it)
{
    sc2::UNIT_TYPEID id = string_to_id(it.key());
    std::string name = sc2::UnitTypeToName(id);
    std::vector<BuildAlternative> build_alternatives;

    if (it.value().find("builds") != it.value().end())
    {
        auto & builds = it.value()["builds"];
        for (auto & build_item : builds)
        {
            BuildAlternative build_alternative = parse_build_alternative(build_item);
            build_alternatives.push_back(build_alternative);
        }
    }
    return { id, build_alternatives };
}

void TechTreeImproved::LoadData() {
    // TODO: Do not hardcode this. Use the latest json available.
    std::ifstream i("techtree.json");
    json j;
    i >> j;

    for (auto & race : j)
    {
        for (json::iterator it = race.begin(); it != race.end(); ++it)
        {
            auto pair = parse_unit(it);
            data[pair.first] = pair.second;
        }
    }

    std::cout << "Done" << std::endl;
}

const std::vector<BuildAlternative> & TechTreeImproved::GetBuildAlternatives(sc2::UNIT_TYPEID unit) const
{
    if (data.count(unit) > 0)
    {
        return data.at(unit);
    }
    else
    {
        std::cout << "No information about unit type " << sc2::UnitTypeToName(unit) << " (" << static_cast<int>(unit) << ")" << std::endl;
        return {};
    }
}

std::vector<std::pair<sc2::UNIT_TYPEID, BuildAlternative>> TechTreeImproved::HowToBuild(sc2::UNIT_TYPEID unit_type)
{
    std::vector<std::pair<sc2::UNIT_TYPEID, BuildAlternative>> alternatives;
    for (auto & pair : data)
    {
        for (const BuildAlternative & alternative : pair.second)
        {
            if (alternative.type == unit_type)
            {
                alternatives.push_back({ pair.first, alternative });
            }
        }
    }
    return alternatives;
}