package com.engebretson;

import org.openjdk.jmh.annotations.*;
import org.openjdk.jmh.infra.Blackhole;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * Comprehensive benchmark testing polymorphic call site performance across collection operations.
 * Tests constructors, putAll/addAll, and utility methods to identify similar patterns to HashMap(Map).
 */
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
@State(Scope.Benchmark)
@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
@Measurement(iterations = 5, time = 3, timeUnit = TimeUnit.SECONDS)
@Fork(value = 1, jvmArgs = {"-XX:+UseParallelGC", "-Xmx3g"})
public class CollectionPolymorphismBenchmark {

    private static final int POISON_ITERATIONS = 10000;
    private static final double CAPACITY_FACTOR = 1.35;

    @Param({"0", "1", "5", "15", "75"})
    private int mapSize;

    private HashMap<String, Integer> inputHashMap;
    private TreeMap<String, Integer> inputTreeMap;
    private LinkedHashMap<String, Integer> inputLinkedHashMap;
    
    private HashSet<String> inputHashSet;
    private TreeSet<String> inputTreeSet;
    private LinkedHashSet<String> inputLinkedHashSet;
    
    private ArrayList<String> inputArrayList;
    private LinkedList<String> inputLinkedList;
    
    private Set<String> inputSet;

    @Setup(Level.Trial)
    public void setup() {
        // Create test data
        inputHashMap = new HashMap<>();
        inputTreeMap = new TreeMap<>();
        inputLinkedHashMap = new LinkedHashMap<>();
        inputHashSet = new HashSet<>();
        inputTreeSet = new TreeSet<>();
        inputLinkedHashSet = new LinkedHashSet<>();
        inputArrayList = new ArrayList<>();
        inputLinkedList = new LinkedList<>();
        inputSet = inputHashMap.keySet();

        for (int i = 0; i < mapSize; i++) {
            String key = "key" + i;
            Integer value = i;
            inputHashMap.put(key, value);
            inputTreeMap.put(key, value);
            inputLinkedHashMap.put(key, value);
            inputHashSet.add(key);
            inputTreeSet.add(key);
            inputLinkedHashSet.add(key);
            inputArrayList.add(key);
            inputLinkedList.add(key);
        }

        // Poison all call sites with polymorphic calls
        poisonCallSites();
    }

    private void poisonCallSites() {
        for (int i = 0; i < POISON_ITERATIONS; i++) {
            Map<String, Integer> sourceMap = (i % 3 == 0) ? inputHashMap : (i % 3 == 1) ? inputTreeMap : inputLinkedHashMap;
            Set<String> sourceSet = (i % 3 == 0) ? inputHashSet : (i % 3 == 1) ? inputTreeSet : inputLinkedHashSet;
            List<String> sourceList = (i % 2 == 0) ? inputArrayList : inputLinkedList;
            
            // Poison constructors
            new HashMap<>(sourceMap);
            new TreeMap<>(sourceMap);
            new LinkedHashMap<>(sourceMap);
            new HashSet<>(sourceSet);
            new TreeSet<>(sourceSet);
            new LinkedHashSet<>(sourceSet);
            new ArrayList<>(sourceList);
            new LinkedList<>(sourceList);
            
            // Poison putAll/addAll
            HashMap<String, Integer> tempMap = new HashMap<>();
            tempMap.putAll(sourceMap);
            HashSet<String> tempSet = new HashSet<>();
            tempSet.addAll(inputSet);
            ArrayList<String> tempList = new ArrayList<>();
            tempList.addAll(sourceList);
            
            // Poison Collections utility methods
            HashMap<String, Integer> destMap = new HashMap<>();
            for (Map.Entry<String, Integer> entry : sourceMap.entrySet()) {
                destMap.put(entry.getKey(), entry.getValue());
            }
        }
    }

    // TreeMap Constructor Tests
    @Benchmark
    public TreeMap<String, Integer> treeMapConstructor_baseline(Blackhole bh) {
        TreeMap<String, Integer> result = new TreeMap<>(inputHashMap);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public TreeMap<String, Integer> treeMapConstructor_manual(Blackhole bh) {
        TreeMap<String, Integer> result = new TreeMap<>();
        for (Map.Entry<String, Integer> entry : inputHashMap.entrySet()) {
            result.put(entry.getKey(), entry.getValue());
        }
        bh.consume(result);
        return result;
    }

    // LinkedHashMap Constructor Tests
    @Benchmark
    public LinkedHashMap<String, Integer> linkedHashMapConstructor_baseline(Blackhole bh) {
        LinkedHashMap<String, Integer> result = new LinkedHashMap<>(inputHashMap);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public LinkedHashMap<String, Integer> linkedHashMapConstructor_manual(Blackhole bh) {
        LinkedHashMap<String, Integer> result = new LinkedHashMap<>((int) (inputHashMap.size() * CAPACITY_FACTOR));
        for (Map.Entry<String, Integer> entry : inputHashMap.entrySet()) {
            result.put(entry.getKey(), entry.getValue());
        }
        bh.consume(result);
        return result;
    }

    // HashSet Constructor Tests
    @Benchmark
    public HashSet<String> hashSetConstructor_baseline(Blackhole bh) {
        HashSet<String> result = new HashSet<>(inputHashSet);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public HashSet<String> hashSetConstructor_manual(Blackhole bh) {
        HashSet<String> result = new HashSet<>((int) (inputHashSet.size() * CAPACITY_FACTOR));
        for (String item : inputHashSet) {
            result.add(item);
        }
        bh.consume(result);
        return result;
    }

    // putAll Tests
    @Benchmark
    public HashMap<String, Integer> putAll_baseline(Blackhole bh) {
        HashMap<String, Integer> result = new HashMap<>((int) (inputHashMap.size() * CAPACITY_FACTOR));
        result.putAll(inputHashMap);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public HashMap<String, Integer> putAll_manual(Blackhole bh) {
        HashMap<String, Integer> result = new HashMap<>((int) (inputHashMap.size() * CAPACITY_FACTOR));
        for (Map.Entry<String, Integer> entry : inputHashMap.entrySet()) {
            result.put(entry.getKey(), entry.getValue());
        }
        bh.consume(result);
        return result;
    }

    // addAll Tests
    @Benchmark
    public HashSet<String> addAll_baseline(Blackhole bh) {
        HashSet<String> result = new HashSet<>((int) (inputSet.size() * CAPACITY_FACTOR));
        result.addAll(inputSet);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public HashSet<String> addAll_manual(Blackhole bh) {
        HashSet<String> result = new HashSet<>((int) (inputSet.size() * CAPACITY_FACTOR));
        for (String key : inputSet) {
            result.add(key);
        }
        bh.consume(result);
        return result;
    }

    // ArrayList Constructor Tests
    @Benchmark
    public ArrayList<String> arrayListConstructor_baseline(Blackhole bh) {
        ArrayList<String> result = new ArrayList<>(inputSet);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public ArrayList<String> arrayListConstructor_manual(Blackhole bh) {
        ArrayList<String> result = new ArrayList<>(inputSet.size());
        for (String item : inputSet) {
            result.add(item);
        }
        bh.consume(result);
        return result;
    }

    // LinkedList Constructor Tests
    @Benchmark
    public LinkedList<String> linkedListConstructor_baseline(Blackhole bh) {
        LinkedList<String> result = new LinkedList<>(inputSet);
        bh.consume(result);
        return result;
    }

    @Benchmark
    public LinkedList<String> linkedListConstructor_manual(Blackhole bh) {
        LinkedList<String> result = new LinkedList<>();
        for (String item : inputSet) {
            result.add(item);
        }
        bh.consume(result);
        return result;
    }
}