package top.r3944realms.lib39.core.sync;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.Entity;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * The type Sync data 2 manager.
 */
@SuppressWarnings({"unused", "DuplicatedCode"})
public class SyncData2Manager {
    private final Map<ResourceLocation, TypedSyncEntry<?, ?>> typedEntries = Maps.newConcurrentMap();

    /**
     * 数据提供者接口 - 用于通过键获取数据
     *
     * @param <K> the type parameter
     * @param <T> the type parameter
     */
    @FunctionalInterface
    public interface DataProvider<K, T> {
        /**
         * 通过键获取数据的 Optional
         *
         * @param key 键
         * @return 数据的 Optional
         */
        Optional<T> getData(K key);
    }

    private static class TypedSyncEntry<K, T extends ISyncData<?>> {
        /**
         * The Manager.
         */
        final ISyncManager<K, T> manager;
        /**
         * The Data provider.
         */
        @Nullable
        final DataProvider<Entity, T> dataProvider;
        /**
         * The Allowed classes.
         */
        final Set<Class<?>> allowedClasses;

        /**
         * Instantiates a new Typed sync entry.
         *
         * @param manager      the manager
         * @param dataProvider the data provider
         */
        TypedSyncEntry(ISyncManager<K, T> manager, @Nullable DataProvider<Entity, T> dataProvider) {
            this.manager = manager;
            this.dataProvider = dataProvider;
            this.allowedClasses = Sets.newConcurrentHashSet();
        }
    }

    /**
     * Register manager with data provider.
     *
     * @param <K>          the type parameter
     * @param <T>          the type parameter
     * @param key          the key
     * @param manager      the manager
     * @param dataProvider the data provider
     */
    public <K, T extends ISyncData<?>> void registerManagerWithProvider(
            ResourceLocation key,
            ISyncManager<K, T> manager,
            DataProvider<Entity, T> dataProvider
    ) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(manager, "Sync manager cannot be null");
        Objects.requireNonNull(dataProvider, "Data provider cannot be null");

        typedEntries.put(key, new TypedSyncEntry<>(manager, dataProvider));
    }

    /**
     * Register manager with function getter.
     *
     * @param <K>     the type parameter
     * @param <T>     the type parameter
     * @param key     the key
     * @param manager the manager
     * @param getter  the data getter function
     */
    public <K, T extends ISyncData<?>> void registerManager(
            ResourceLocation key,
            ISyncManager<K, T> manager,
            Function<Entity, Optional<T>> getter
    ) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(manager, "Sync manager cannot be null");
        Objects.requireNonNull(getter, "Data getter function cannot be null");

        typedEntries.put(key, new TypedSyncEntry<>(manager, getter::apply));
    }

    /**
     * 向后兼容的注册方法（只注册管理器，不注册数据提供者）
     *
     * @param key     the key
     * @param manager the manager
     */
    @SuppressWarnings("unchecked")
    public void registerManager(ResourceLocation key, ISyncManager<?, ? extends ISyncData<?>> manager) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(manager, "Sync manager cannot be null");

        // 创建一个没有数据提供者的 TypedSyncEntry
        typedEntries.put(key, new TypedSyncEntry<>(
                (ISyncManager<?, ISyncData<?>>) manager,
                null
        ));
    }

    /**
     * Gets manager.
     *
     * @param <K> the type parameter
     * @param <T> the type parameter
     * @param key the key
     * @return the manager
     */
    @SuppressWarnings("unchecked")
    public <K, T extends ISyncData<?>> Optional<ISyncManager<K, T>> getManager(ResourceLocation key) {
        TypedSyncEntry<?,?> entry = typedEntries.get(key);
        return entry != null ? Optional.of((ISyncManager<K,T>) entry.manager) : Optional.empty();
    }

    /**
     * Gets data provider.
     *
     * @param <T> the type parameter
     * @param key the key
     * @return the data provider
     */
    @SuppressWarnings("unchecked")
    public <T extends ISyncData<?>> Optional<DataProvider<Entity, T>> getDataProvider(ResourceLocation key) {
        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry != null && entry.dataProvider != null) {
            return Optional.of((DataProvider<Entity, T>) entry.dataProvider);
        }
        return Optional.empty();
    }

    /**
     * 获取实体数据
     *
     * @param <T>    the type parameter
     * @param key    the key
     * @param entity the entity
     * @return the entity data
     */
    @SuppressWarnings("unchecked")
    public <T extends ISyncData<?>> Optional<T> getEntityData(ResourceLocation key, Entity entity) {
        return getDataProvider(key)
                .flatMap(provider -> {
                    Optional<ISyncData<?>> result = provider.getData(entity);
                    return (Optional<T>) result;
                });
    }

    /**
     * Allow entity class.
     *
     * @param key     the key
     * @param classes the classes
     */
    public final void allowEntityClass(ResourceLocation key, Class<?>... classes) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(classes, "Classes array cannot be null");

        if (classes.length == 0) {
            return;
        }

        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry != null) {
            entry.allowedClasses.addAll(Arrays.asList(classes));
        }
    }

    /**
     * 移除允许的实体类
     *
     * @param key     the key
     * @param classes the classes
     */
    public final void disallowEntityClass(ResourceLocation key, Class<?>... classes) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(classes, "Classes array cannot be null");

        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry != null && classes.length > 0) {
            Arrays.asList(classes).forEach(entry.allowedClasses::remove);
        }
    }

    /**
     * 绑定数据提供者（用于分离注册的情况）
     *
     * @param <T>          the type parameter
     * @param key          the key
     * @param dataProvider the data provider
     */
    public <T extends ISyncData<?>> void bindDataProvider(ResourceLocation key, DataProvider<Entity, T> dataProvider) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(dataProvider, "Data provider cannot be null");

        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry != null) {
            // 更新现有条目的数据提供者
            updateDataProviderInEntry(key, entry, dataProvider);
        } else {
            throw new IllegalArgumentException("No manager found for " + key);
        }
    }

    /**
     * 绑定简单的数据获取器
     *
     * @param <T>    the type parameter
     * @param key    the key
     * @param getter the data getter function
     */
    public <T extends ISyncData<?>> void bindDataGetter(ResourceLocation key, Function<Entity, Optional<T>> getter) {
        bindDataProvider(key, getter::apply);
    }

    /**
     * 解绑数据提供者
     *
     * @param key the key
     */
    public void unbindDataProvider(ResourceLocation key) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");

        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry != null) {
            // 将数据提供者设置为null，但保留管理器和其他配置
            updateDataProviderInEntry(key, entry, null);
        }
    }

    /**
     * 清除允许的实体类
     *
     * @param key the key
     */
    public void clearAllowedEntityClasses(ResourceLocation key) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");

        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry != null) {
            entry.allowedClasses.clear();
        }
    }

    /**
     * Is entity class allowed boolean.
     *
     * @param key         the key
     * @param entityClass the entity class
     * @return the boolean
     */
    public boolean isEntityClassAllowed(ResourceLocation key, Class<?> entityClass) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        Objects.requireNonNull(entityClass, "Entity class cannot be null");

        TypedSyncEntry<?, ?> entry = typedEntries.get(key);
        if (entry == null) {
            return false;
        }

        // 如果没有设置允许的类，则默认允许所有类
        if (entry.allowedClasses.isEmpty()) {
            return true;
        }

        // 检查实体类是否在允许的类中
        return entry.allowedClasses.stream()
                .anyMatch(allowedClass -> allowedClass.isAssignableFrom(entityClass));
    }

    /**
     * Track entity for manager.
     *
     * @param entity    the entity
     * @param managerId the manager id
     */
    @SuppressWarnings("unchecked")
    public void trackEntityForManager(Entity entity, ResourceLocation managerId) {
        TypedSyncEntry<UUID, ?> entry = (TypedSyncEntry<UUID, ?>) typedEntries.get(managerId);
        if (entry != null) {
            trackEntityWithTypedEntry(entity, entry);
        }
    }

    private <T extends ISyncData<?>> void trackEntityWithTypedEntry(Entity entity, @NotNull TypedSyncEntry<UUID, T> entry) {
        if (entry.dataProvider != null) {
            entry.dataProvider.getData(entity)
                    .ifPresent(data -> entry.manager.track(entity.getUUID(), data));
        }
    }

    /**
     * Untrack entity for manager.
     *
     * @param entity    the entity
     * @param managerId the manager id
     */
    @SuppressWarnings("unchecked")
    public void untrackEntityForManager(Entity entity, ResourceLocation managerId) {
        TypedSyncEntry<UUID, ?> entry = (TypedSyncEntry<UUID, ?>) typedEntries.get(managerId);
        if (entry != null) {
            untrackEntityWithTypedEntry(entity, entry);
        }
    }

    private <T extends ISyncData<?>> void untrackEntityWithTypedEntry(Entity entity, @NotNull TypedSyncEntry<UUID, T> entry) {
        if (entry.dataProvider != null) {
            entry.dataProvider.getData(entity)
                    .ifPresent(data -> entry.manager.untrack(entity.getUUID(), data));
        }
    }

    /**
     * 从所有管理器中移除实体跟踪
     *
     * @param entity the entity
     */
    public void untrackEntityFromAllManagers(Entity entity) {
        for (ResourceLocation id : getRegisteredKeys()) {
            if (isEntityClassAllowed(id, entity.getClass())) {
                untrackEntityForManager(entity, id);
            }
        }
    }

    /**
     * 批量从管理器中移除实体跟踪
     *
     * @param entities  the entities
     * @param managerId the manager id
     */
    public void untrackEntitiesForManager(@NotNull Iterable<Entity> entities, ResourceLocation managerId) {
        for (Entity entity : entities) {
            untrackEntityForManager(entity, managerId);
        }
    }

    /**
     * 从所有管理器中批量移除实体跟踪
     *
     * @param entities the entities
     */
    public void untrackEntitiesFromAllManagers(@NotNull Iterable<Entity> entities) {
        for (Entity entity : entities) {
            untrackEntityFromAllManagers(entity);
        }
    }

    /**
     * 强制清理管理器中的所有跟踪数据
     *
     * @param managerId the manager id
     */
    public void clearAllTrackedData(ResourceLocation managerId) {
        TypedSyncEntry<?, ?> entry = typedEntries.get(managerId);
        if (entry != null) {
            clearTrackedDataForEntry(entry);
        }
    }

    private <K, T extends ISyncData<?>> void clearTrackedDataForEntry(@NotNull TypedSyncEntry<K, T> entry) {
        Set<T> syncSet = entry.manager.getSyncSet();
        if (syncSet != null) {
            syncSet.clear();
        }
    }

    /**
     * 清理所有管理器的跟踪数据
     */
    public void clearAllTrackedData() {
        for (ResourceLocation id : getRegisteredKeys()) {
            clearAllTrackedData(id);
        }
    }

    // 辅助方法：更新条目的数据提供者
    @SuppressWarnings("unchecked")
    private <K, T extends ISyncData<?>> void updateDataProviderInEntry(
            ResourceLocation id,
            TypedSyncEntry<?,?> entry,
            DataProvider<Entity, T> newDataProvider
    ) {
        // 由于 DataProvider 是 final，我们需要创建一个新的 TypedSyncEntry
        TypedSyncEntry<K, T> newEntry = new TypedSyncEntry<>(
                (ISyncManager<K, T>) entry.manager,
                newDataProvider
        );
        newEntry.allowedClasses.addAll(entry.allowedClasses);

        typedEntries.put(id, newEntry);
    }

    /**
     * Gets registered keys.
     *
     * @return the registered keys
     */
    public Set<ResourceLocation> getRegisteredKeys() {
        return Collections.unmodifiableSet(typedEntries.keySet());
    }

    /**
     * For each.
     *
     * @param consumer the consumer
     */
    public void forEach(BiConsumer<ResourceLocation, ISyncManager<?,?>> consumer) {
        Objects.requireNonNull(consumer, "Consumer cannot be null");
        typedEntries.forEach((key, entry) -> consumer.accept(key, entry.manager));
    }

    /**
     * Gets manager count.
     *
     * @return the manager count
     */
    public int getManagerCount() {
        return typedEntries.size();
    }

    /**
     * Clear all.
     */
    public void clearAll() {
        typedEntries.clear();
    }

    /**
     * 移除管理器（包括所有相关配置）
     *
     * @param key the key
     */
    public void removeManager(ResourceLocation key) {
        Objects.requireNonNull(key, "ResourceLocation key cannot be null");
        typedEntries.remove(key);
    }
}