#include "codegen.h"
#include "analysis.h"
#include <map>
#include <string>
#include <variant>

std::map<std::string, std::string> CATEGORY_STR = {
    { "APPEARANCE", "PROP_CATEGORY_APPEARENCE" },
    { "DATA", "PROP_CATEGORY_DATA" },
    { "BEHAVIOR", "PROP_CATEGORY_BEHAVIOR" },
    { "PART", "PROP_CATEGORY_PART" },
    { "SURFACE", "PROP_CATEGORY_SURFACE" },
};

std::map<std::string, std::string> MAPPED_TYPE = {
    { "bool", "Data::Bool" },
    { "int", "Data::Int" },
    { "float", "Data::Float" },
    { "std::string", "Data::String" },
    { "glm::vec3", "Vector3" },
};

std::map<std::string, std::monostate> ENUM_TYPES = {
    { "SurfaceType", std::monostate() }
};

std::string parseWeakPtr(std::string weakPtrType) {
    if (!weakPtrType.starts_with("std::weak_ptr")) return "";

    int pos0 = weakPtrType.find("<");
    int pos1 = weakPtrType.find(">");

    std::string subtype = weakPtrType.substr(pos0+1, pos1-pos0-1);
    return subtype;
}

std::string castFromVariant(std::string valueStr, std::string fieldType) {
    // Manual exception for now, enums will get their own system eventually
    if (fieldType == "SurfaceType") {
        return "(SurfaceType)(int)" + valueStr + ".get<Data::Int>()";
    }

    std::string mappedType = MAPPED_TYPE[fieldType];
    return valueStr + ".get<" + (!mappedType.empty() ? mappedType : fieldType) + ">()";
}

std::string castToVariant(std::string valueStr, std::string fieldType) {
    // Manual exception for now, enums will get their own system eventually
    if (fieldType == "SurfaceType") {
        return "Data::Int((int)" + valueStr + ")";
    }

    // InstanceRef
    std::string subtype = parseWeakPtr(fieldType);
    if (!subtype.empty()) {
        return "Data::Variant(" + valueStr + ".expired() ? Data::InstanceRef() : Data::InstanceRef(std::dynamic_pointer_cast<Instance>(" + valueStr + ".lock())))";
    }

    std::string mappedType = MAPPED_TYPE[fieldType];
    if (!mappedType.empty()) {
        return mappedType + "(" + valueStr + ")";
    }
    return valueStr;
}

void writePropertySetHandler(std::ofstream& out, ClassAnalysis state) {
    out << "fallible<MemberNotFound, AssignToReadOnlyMember> " << state.name << "::InternalSetPropertyValue(std::string name, Data::Variant value) {";

    out << "\n    ";
    bool first = true;
    for (auto& prop : state.properties) {
        out << (first ? "" : " else ") << "if (name == \"" << prop.name << "\") {";
        // InstanceRef
        std::string subtype = parseWeakPtr(prop.backingFieldType);

        if (prop.flags & PropertyFlag_Readonly) {
            out << "\n        return AssignToReadOnlyMember(GetClass()->className, name)";
        } else if (prop.cframeMember == CFrameMember_Position) {
            out << "\n        this->" << prop.fieldName << " = this->" << prop.fieldName << ".Rotation() + value.get<Vector3>();";
        }  else if (prop.cframeMember == CFrameMember_Rotation) {
            out << "\n        this->" << prop.fieldName << " = CFrame::FromEulerAnglesXYZ(value.get<Vector3>()) + this->" << prop.fieldName << ".Position();";
        } else if (!subtype.empty()) {
            out << "\n        std::weak_ptr<Instance> ref = value.get<Data::InstanceRef>();"
                << "\n        this->" << prop.fieldName << " = ref.expired() ? std::weak_ptr<" << subtype << ">() : std::dynamic_pointer_cast<" << subtype << ">(ref.lock());";
        } else {
            out << "\n        this->" << prop.fieldName << " = " << castFromVariant("value", prop.backingFieldType) << ";";
        }

        out << "\n    }";
        first = false;
    }

    // If it's empty, just return the parent's impl
    if (state.properties.empty()) {
        out << "\n    return " << state.baseClass << "::InternalSetPropertyValue(name, value);";
    } else {
        // Otherwise, add else and return
        out << "else {";
        out << "\n        return " << state.baseClass << "::InternalSetPropertyValue(name, value);";
        out << "\n    }";
        out << "\n    return {};";
    }

    out << "\n};\n\n";
}

void writePropertyUpdateHandler(std::ofstream& out, ClassAnalysis state) {
    out << "void " << state.name << "::InternalUpdateProperty(std::string name) {";

    out << "\n    ";
    bool first = true;
    for (auto& prop : state.properties) {
        if (prop.onUpdateCallback.empty()) continue;
        out << (first ? "" : " else ") << "if (name == \"" << prop.name << "\") {";

        out << "\n        " << prop.onUpdateCallback << "(name);";

        out << "\n    }";
        first = false;
    }

    out << "\n    " << state.baseClass << "::InternalUpdateProperty(name);";

    out << "\n};\n\n";
}

void writePropertyGetHandler(std::ofstream& out, ClassAnalysis state) {
    out << "result<Data::Variant, MemberNotFound> " << state.name << "::InternalGetPropertyValue(std::string name) {";

    out << "\n    ";
    bool first = true;
    for (auto& prop : state.properties) {
        out << (first ? "" : " else ") << "if (name == \"" << prop.name << "\") {";

        if (prop.cframeMember == CFrameMember_Position) {
            out << "\n        return Data::Variant(" << prop.fieldName << ".Position());";
        } else if (prop.cframeMember == CFrameMember_Rotation) {
            out << "\n        return Data::Variant(" << prop.fieldName << ".ToEulerAnglesXYZ());";
        } else {
            out << "\n        return Data::Variant(" << castToVariant(prop.fieldName, prop.backingFieldType) << ");";
        }

        out << "\n    }";
        first = false;
    }
    
    out << "\n    return " << state.baseClass << "::InternalGetPropertyValue(name);";
    // out << "\n    return MemberNotFound(\"" << state.name << "\", name);";

    out << "\n};\n\n";
}

void writePropertiesList(std::ofstream& out, ClassAnalysis state) {
    out << "std::vector<std::string> " << state.name << "::InternalGetProperties() {\n";
    out << "    std::vector<std::string> properties = " << state.baseClass << "::InternalGetProperties();\n";

    for (auto& prop : state.properties) {
        out << "    properties.push_back(\"" << prop.name << "\");\n";
    }

    out << "    return properties;\n";

    out << "};\n\n";
}

void writePropertyMetaHandler(std::ofstream& out, ClassAnalysis state) {
    out << "result<PropertyMeta, MemberNotFound> " << state.name << "::InternalGetPropertyMeta(std::string name) {";

    out << "\n    ";
    bool first = true;
    for (auto& prop : state.properties) {
        out << (first ? "" : " else ") << "if (name == \"" << prop.name << "\") {";

        std::string type = MAPPED_TYPE[prop.backingFieldType];
        if (type.empty()) type = prop.backingFieldType;
        if (type == "SurfaceType") type = "Data::Int";
        if (!parseWeakPtr(prop.backingFieldType).empty()) type = "Data::InstanceRef";

        std::string strFlags;
        if (prop.flags & PropertyFlag_Readonly)
            strFlags += " | PROP_READONLY";
        if (prop.flags & PropertyFlag_Hidden)
            strFlags += " | PROP_HIDDEN";
        if (prop.flags & PropertyFlag_NoSave)
            strFlags += " | PROP_NOSAVE";
        if (prop.flags & PropertyFlag_UnitFloat)
            strFlags += " | PROP_UNIT_FLOAT";
        if (!strFlags.empty()) strFlags = strFlags.substr(3); // Remove leading pipe
        else strFlags = "0"; // 0 == No option

        std::string category = CATEGORY_STR[prop.category];
        if (category.empty()) category = "PROP_CATEGORY_DATA";

        out << "\n        return PropertyMeta { &" << type << "::TYPE, " << strFlags << ", " << category << " };";

        out << "\n    }";
        first = false;
    }
    
    out << "\n    return " << state.baseClass << "::InternalGetPropertyMeta(name);";
    // out << "\n    return MemberNotFound(\"" << state.name << "\", name);";

    out << "\n};\n\n";
}

void writeCodeForClass(std::ofstream& out, ClassAnalysis& state) {
    std::string strFlags;
    if (state.flags & ClassFlag_NotCreatable)
        strFlags += " | INSTANCE_NOTCREATABLE";
    if (state.flags & ClassFlag_Service)
        strFlags += " | INSTANCE_SERVICE";
    if (state.flags & ClassFlag_Hidden)
        strFlags += " | INSTANCE_HIDDEN";
    if (!strFlags.empty()) strFlags = strFlags.substr(3); // Remove leading pipe
    else strFlags = "0"; // 0 == No option

    out << "/////////////////////////////////////////////////////////////////////////////////////////\n";
    out << "// This file was automatically generated by autogen, and should not be edited manually //\n";
    out << "/////////////////////////////////////////////////////////////////////////////////////////\n\n";

    std::string constructorStr;
    if (state.abstract) constructorStr = "nullptr";
    else constructorStr = "&" + state.name + "::Create";

    out << "#define __AUTOGEN_EXTRA_INCLUDES__\n";
    out << "#include \"" << state.headerPath << "\"\n\n";
    out << "const InstanceType " << state.name << "::TYPE = {\n"
        << "    .super = &" << state.baseClass << "::TYPE,\n"
        << "    .className = \"" << state.name << "\",\n"
        << "    .constructor = " << constructorStr << ",\n"
        << "    .explorerIcon = \"" << state.explorerIcon << "\",\n"
        << "    .flags = " << strFlags << ",\n"
        << "};\n\n";

    out << "const InstanceType* " << state.name << "::GetClass() {\n"
        << "    return &TYPE;\n"
        << "};\n\n";

    // Special case for our Service class
    if (state.baseClass == "Service") state.baseClass = "Instance";

    writePropertySetHandler(out, state);
    writePropertyGetHandler(out, state);
    writePropertyMetaHandler(out, state);
    writePropertyUpdateHandler(out, state);
    writePropertiesList(out, state);
}