package org.shedlang.compiler.util;

import org.pcollections.Empty;
import org.pcollections.PMap;
import org.pcollections.PVector;

import java.util.Optional;
import java.util.stream.Stream;

public class POrderedMap<K, V> {
    public static <K, V> POrderedMap<K, V> empty() {
        return new POrderedMap<>(Empty.map(), Empty.vector());
    }

    private final PMap<K, V> map;
    private final PVector<K> keys;

    private POrderedMap(
        PMap<K, V> map,
        PVector<K> keys
    ) {
        this.map = map;
        this.keys = keys;
    }

    public POrderedMap<K, V> plus(K key, V value) {
        return new POrderedMap<>(map.plus(key, value), keys.plus(key));
    }

    public POrderedMap<K, V> plusAll(POrderedMap<K, V> other) {
        var newMap = this;
        for (var otherKey : other.keys) {
            newMap = newMap.plus(otherKey, other.map.get(otherKey));
        }
        return newMap;
    }

    public boolean containsKey(K key) {
        return map.containsKey(key);
    }

    public PVector<K> keys() {
        return keys;
    }

    public Optional<V> get(K key) {
        return Optional.ofNullable(map.get(key));
    }

    public int size() {
        return keys.size();
    }

    public Stream<V> valueStream() {
        return keys.stream()
            .map(key -> map.get(key));
    }

    @Override
    public boolean equals(Object other) {
        if (this == other) {
            return true;
        }
        if (other == null) {
            return false;
        }
        if (!(other instanceof POrderedMap<?, ?> otherMap)) {
            return false;
        }
        return map.equals(otherMap.map) && keys.equals(otherMap.keys);
    }

    @Override
    public int hashCode() {
        int result = map.hashCode();
        result = 31 * result + keys.hashCode();
        return result;
    }
}