#include "UnistylesRegistry.h"
#include "UnistylesState.h"
#include "Parser.h"

using namespace margelo::nitro::unistyles;
using namespace facebook;
using namespace facebook::react;

std::atomic<int> core::UnistylesRegistry::_nextStyleSheetTag{0};

void core::UnistylesRegistry::registerTheme(jsi::Runtime& rt, std::string name, jsi::Value& theme) {
    auto& state = this->getState();

    state._jsThemes.emplace(name, std::move(theme));
    state._registeredThemeNames.push_back(name);
}

void core::UnistylesRegistry::registerBreakpoints(std::vector<std::pair<std::string, double>>& sortedBreakpoints) {
    auto& state = this->getState();

    state._sortedBreakpointPairs = std::move(sortedBreakpoints);
}

void core::UnistylesRegistry::setPrefersAdaptiveThemes(bool prefersAdaptiveThemes) {
    auto& state = this->getState();

    state._prefersAdaptiveThemes = prefersAdaptiveThemes;
}

void core::UnistylesRegistry::setInitialThemeName(std::string themeName) {
    auto& state = this->getState();

    state._initialThemeName = themeName;
}

core::UnistylesState& core::UnistylesRegistry::getState() {
    if (!this->_state) {
        throw std::runtime_error("Unistyles was loaded, but it's not configured. Did you forget to call StyleSheet.configure? If you don't want to use any themes or breakpoints, simply call it with an empty object {}.");
    }

    return *this->_state;
}

void core::UnistylesRegistry::createState() {
    this->_state = std::make_unique<UnistylesState>();
}

void core::UnistylesRegistry::updateTheme(jsi::Runtime& rt, std::string& themeName, jsi::Function&& callback) {
    auto& state = this->getState();
    auto it = state._jsThemes.find(themeName);

    helpers::assertThat(rt, it != state._jsThemes.end(), "Unistyles: You're trying to update theme '" + themeName + "' but it wasn't registered.");

    auto result = callback.call(rt, it->second);

    helpers::assertThat(rt, result.isObject(), "Unistyles: Returned theme is not an object. Please check your updateTheme function.");

    it->second = result.asObject(rt);
}

void core::UnistylesRegistry::linkShadowNodeWithUnistyle(
    jsi::Runtime& rt,
    const ShadowNodeFamily* shadowNodeFamily,
    std::vector<std::shared_ptr<UnistyleData>>& unistylesData,
    std::optional<folly::dynamic> initialScopedUpdate
) {
    this->trafficController.withLock([this, &rt, &unistylesData, shadowNodeFamily, &initialScopedUpdate](){
        // Clear suspension state if this family was previously suspended
        if (_suspendedFamilies.erase(shadowNodeFamily) > 0) {
            auto* mutableFamily = const_cast<ShadowNodeFamily*>(shadowNodeFamily);
            mutableFamily->nativeProps_DEPRECATED.reset();
            // Clear old registry entries to prevent stale UnistyleData accumulation
            this->_shadowRegistry.erase(shadowNodeFamily);
            // Remove any stale traffic controller entry (e.g. from a theme change during suspension)
            this->trafficController.removeShadowNode(shadowNodeFamily);
        }

        std::for_each(unistylesData.begin(), unistylesData.end(), [this, shadowNodeFamily](std::shared_ptr<UnistyleData> unistyleData){
            this->_shadowRegistry[shadowNodeFamily].emplace_back(unistyleData);
        });

        // Required for scoped themes to apply on initial mount
        if (initialScopedUpdate.has_value()) {
            shadow::ShadowLeafUpdates updates;

            updates.emplace(shadowNodeFamily, std::move(*initialScopedUpdate));
            this->trafficController.setUpdates(updates);
        }
    });
}

void core::UnistylesRegistry::removeDuplicatedUnistyles(const ShadowNodeFamily *shadowNodeFamily, std::vector<core::Unistyle::Shared>& unistyles) {
    auto targetFamilyUnistyles = this->_shadowRegistry[shadowNodeFamily];

    unistyles.erase(
        std::remove_if(
            unistyles.begin(),
            unistyles.end(),
            [&targetFamilyUnistyles](const core::Unistyle::Shared& unistyle) {
                return std::any_of(
                    targetFamilyUnistyles.begin(),
                    targetFamilyUnistyles.end(),
                    [&unistyle](const std::shared_ptr<core::UnistyleData>& data) {
                        return data->unistyle == unistyle;
                    }
                );
            }
        ),
        unistyles.end()
    );
}

void core::UnistylesRegistry::unlinkShadowNodeWithUnistyles(const ShadowNodeFamily* shadowNodeFamily) {
    this->trafficController.withLock([this, shadowNodeFamily](){
        this->_shadowRegistry.erase(shadowNodeFamily);
        this->_suspendedFamilies.erase(shadowNodeFamily);
        this->trafficController.removeShadowNode(shadowNodeFamily);
    });
}

void core::UnistylesRegistry::suspendShadowNode(const ShadowNodeFamily* shadowNodeFamily) {
    this->trafficController.withLock([this, shadowNodeFamily](){
        if (this->_shadowRegistry.contains(shadowNodeFamily)) {
            this->_suspendedFamilies.insert(shadowNodeFamily);
        }
    });
}

bool core::UnistylesRegistry::isSuspended(const ShadowNodeFamily* family) const noexcept {
    return _suspendedFamilies.count(family) > 0;
}

std::shared_ptr<core::StyleSheet> core::UnistylesRegistry::addStyleSheet(jsi::Runtime& rt, core::StyleSheetType type, jsi::Object&& rawValue) {
    int tag = _nextStyleSheetTag.fetch_add(1);

    auto sheet = std::make_shared<core::StyleSheet>(tag, type, std::move(rawValue));
    this->_styleSheetRegistry[tag] = sheet;

    return sheet;
}

core::DependencyMap core::UnistylesRegistry::buildDependencyMap(std::vector<UnistyleDependency>& deps) {
    core::DependencyMap dependencyMap;

    std::unordered_set<UnistyleDependency> uniqueDependencies(deps.begin(), deps.end());

    for (const auto& [family, unistyles] : this->_shadowRegistry) {
        bool hasAnyOfDependencies = false;

        // Check if any dependency matches
        for (const auto& unistyleData : unistyles) {
            for (const auto& dep : unistyleData->unistyle->dependencies) {
                if (uniqueDependencies.count(dep)) {
                    hasAnyOfDependencies = true;
                    break;
                }
            }

            if (hasAnyOfDependencies) {
                break;
            };
        }

        if (!hasAnyOfDependencies) {
            continue;
        }

        dependencyMap[family].insert(
            dependencyMap[family].end(),
            unistyles.begin(),
            unistyles.end()
        );
    }

    return dependencyMap;
}

// called from proxied function only, we don't know host
// so we need to rebuild all instances as they may have different variants
void core::UnistylesRegistry::shadowLeafUpdateFromUnistyle(jsi::Runtime& rt, Unistyle::Shared unistyle, jsi::Value& maybePressableId) {
    shadow::ShadowLeafUpdates updates;
    this->trafficController.withLock([this, &rt, &maybePressableId, unistyle, &updates](){
        auto parser = parser::Parser(nullptr);
        std::optional<std::string> pressableId = maybePressableId.isString()
            ? std::make_optional(maybePressableId.asString(rt).utf8(rt))
            : std::nullopt;

        for (const auto& [family, unistyles] : this->_shadowRegistry) {
            for (const auto& unistyleData : unistyles) {
                if (unistyleData->unistyle == unistyle) {
                    updates[family] = parser.parseStylesToShadowTreeStyles(rt, { unistyleData });
                }
            }
        }

        this->trafficController.setUpdates(updates);
    });
}

std::vector<std::shared_ptr<core::StyleSheet>>core::UnistylesRegistry::getStyleSheetsToRefresh(std::vector<UnistyleDependency>& unistylesDependencies) {
    std::vector<std::shared_ptr<core::StyleSheet>> stylesheetsToRefresh;
    std::unordered_set<UnistyleDependency> depSet(
        unistylesDependencies.begin(),
        unistylesDependencies.end()
    );

    bool themeDidChange = depSet.count(UnistyleDependency::THEME) > 0;
    bool runtimeDidChange = (themeDidChange && depSet.size() > 1) || !depSet.empty();

    if (!themeDidChange && !runtimeDidChange) {
        return stylesheetsToRefresh;
    }

    auto hasMatchingDependency = [&depSet](const auto& unistyles) {
        for (const auto& [_, unistyle] : unistyles) {
            for (const auto& dep : unistyle->dependencies) {
                if (depSet.count(dep)) {
                    return true;
                }
            }
        }

        return false;
    };

    for (const auto& [_, styleSheet] : this->_styleSheetRegistry) {
        if (styleSheet->type == StyleSheetType::ThemableWithMiniRuntime || styleSheet->type == StyleSheetType::Static) {
            if (hasMatchingDependency(styleSheet->unistyles)) {
                stylesheetsToRefresh.emplace_back(styleSheet);
            }
        }

        if (styleSheet->type == StyleSheetType::Themable && themeDidChange) {
            stylesheetsToRefresh.emplace_back(styleSheet);
        }
    }

    return stylesheetsToRefresh;
}

core::Unistyle::Shared core::UnistylesRegistry::getUnistyleById(std::string unistyleID) {
    for (auto& pair: this->_styleSheetRegistry) {
        auto [_, stylesheet] = pair;

        for (auto unistylePair: stylesheet->unistyles) {
            auto [_, unistyle] = unistylePair;

            if (unistyle->unid == unistyleID) {
                return unistyle;
            }
        }
    }

    return nullptr;
}

const std::optional<std::string> core::UnistylesRegistry::getScopedTheme() {
    return this->_scopedTheme;
}

void core::UnistylesRegistry::setScopedTheme(std::optional<std::string> themeName) {
    this->_scopedTheme = std::move(themeName);
}

void core::UnistylesRegistry::destroy() {
    this->_state.reset();
    this->_styleSheetRegistry.clear();
    this->_shadowRegistry.clear();
    this->_suspendedFamilies.clear();
    this->_scopedTheme = std::nullopt;
    _nextStyleSheetTag.store(0);
}