package org.zwobble.hobgoblin.compiler.analysis;

import org.zwobble.hobgoblin.compiler.analysis.errors.SumTypeHasNoVariantsError;
import org.zwobble.hobgoblin.compiler.analysis.errors.TypeIsInfiniteError;
import org.zwobble.hobgoblin.compiler.builtins.NativeTypes;
import org.zwobble.hobgoblin.compiler.typechecker.TypesInfo;
import org.zwobble.hobgoblin.compiler.types.*;

import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.IntStream;

public class ArbitraryValueAnalysis {
    public sealed interface ArbitraryValue {
        void collectTypeParams(Set<TypeParam> typeParams);

        record NativeFunctionCall(
            SimpleNativeType type,
            List<ArbitraryValue> args
        ) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
                for (var arg : args) {
                    arg.collectTypeParams(typeParams);
                }
            }
        }

        record ConstructedStructBuilder(
            ConstructedStructType type,
            List<ArbitraryValue> args
        ) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
                for (var arg : args) {
                    arg.collectTypeParams(typeParams);
                }
            }
        }

        record SimpleStructBuilder(SimpleStructType type) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
            }
        }

        record EnumVariantValue(EnumType type, EnumVariant variant) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
            }
        }

        record TypeParamThunk(TypeParam typeParam) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
                typeParams.add(typeParam);
            }
        }

        record Box(ArbitraryValue value) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
                value.collectTypeParams(typeParams);
            }
        }

        record Shared(ArbitraryValue value) implements ArbitraryValue {
            @Override
            public void collectTypeParams(Set<TypeParam> typeParams) {
                value.collectTypeParams(typeParams);
            }
        }
    }

    public static ArbitraryValue arbitraryValue(Type type, TypesInfo typesInfo) {
        return switch (type) {
            case ConstructedNativeType constructedNativeType -> {
                if (constructedNativeType.constructor().equals(NativeTypes.BOX)) {
                    yield new ArbitraryValue.Box(arbitraryValue(constructedNativeType.args().getFirst(), typesInfo));
                }

                if (constructedNativeType.constructor().equals(NativeTypes.SHARED)) {
                    yield new ArbitraryValue.Shared(arbitraryValue(constructedNativeType.args().getFirst(), typesInfo));
                }

                var args = constructedNativeType.args().stream()
                    .map(argType -> arbitraryValue(argType, typesInfo))
                    .toList();

                yield new ArbitraryValue.NativeFunctionCall(
                    constructedNativeType.constructor().genericType(),
                    args
                );
            }

            case ConstructedStructType constructedStructType -> {
                var fields = typesInfo.fieldsOf(constructedStructType.constructor().genericType())
                    .orElse(List.of());
                var usedTypeParams = new HashSet<TypeParam>();
                for (var field : fields) {
                    arbitraryValue(field.type(), typesInfo).collectTypeParams(usedTypeParams);
                }

                var allTypeParams = constructedStructType.constructor().params();
                var args = IntStream.range(0, constructedStructType.args().size())
                    .filter(paramIndex -> usedTypeParams.contains(allTypeParams.get(paramIndex)))
                    .mapToObj(paramIndex -> arbitraryValue(constructedStructType.args().get(paramIndex), typesInfo))
                    .toList();

                yield new ArbitraryValue.ConstructedStructBuilder(
                    constructedStructType,
                    args
                );
            }

            case EnumType enumType ->
                new ArbitraryValue.EnumVariantValue(enumType, typesInfo.enumVariants(enumType).getFirst());

            case SimpleNativeType simpleNativeType ->
                new ArbitraryValue.NativeFunctionCall(simpleNativeType, List.of());

            case SimpleStructType simpleStructType ->
                new ArbitraryValue.SimpleStructBuilder(simpleStructType);

            case SumType sumType -> {
                var firstVariant = selectArbitraryVariant(sumType, typesInfo);
                yield arbitraryValue(firstVariant, typesInfo);
            }

            case TypeLevelValueType _ ->
                throw new UnsupportedOperationException("TODO");

            case TypeParam typeParam ->
                new ArbitraryValue.TypeParamThunk(typeParam);
        };
    }

    private static void collectUsedTypeParams(
        ConstructedStructType structType,
        Set<TypeParam> usedTypeParams
    ) {

    }

    public static Type selectArbitraryVariant(SumType sumType, TypesInfo typesInfo) {
        var sumTypeInfo = typesInfo.sumTypeInfo(sumType);
        if (sumTypeInfo.variants().isEmpty()) {
            throw new SumTypeHasNoVariantsError(sumType, sumTypeInfo.source());
        }

        var seenTypes = new HashSet<Type>();
        seenTypes.add(sumType);
        return sumTypeInfo.variants().stream()
            .filter(variant -> isNonRecursivelyConstructable(variant.valueType(), seenTypes, typesInfo))
            .findFirst()
            .orElseThrow(() -> new TypeIsInfiniteError(sumType, sumTypeInfo.source()))
            .containerType();
    }

    private static boolean isNonRecursivelyConstructable(Type type, Set<Type> seenTypes, TypesInfo typesInfo) {
        if (seenTypes.contains(type)) {
            return false;
        }

        return switch (type) {
            case ConstructedType constructedType -> {
                if (
                    constructedType.constructor().equals(NativeTypes.BOX) ||
                    constructedType.constructor().equals(NativeTypes.SHARED)
                ) {
                    yield isNonRecursivelyConstructable(
                        constructedType.args().getFirst(),
                        seenTypes,
                        typesInfo
                    );
                } else if (
                    constructedType.constructor().equals(NativeTypes.LIST) ||
                    constructedType.constructor().equals(NativeTypes.OPTION)
                ) {
                    yield true;
                } else {
                    throw new UnsupportedOperationException("TODO");
                }
            }

            case EnumType _ ->
                true;

            case SimpleNativeType _ ->
                true;

            case SimpleStructType structType -> {
                var newSeenTypes = new HashSet<>(seenTypes);
                newSeenTypes.add(type);
                var fields = typesInfo.fieldsOf(structType);
                yield fields.orElse(List.of()).stream()
                    .allMatch(field -> isNonRecursivelyConstructable(field.type(), newSeenTypes, typesInfo));
            }

            case SumType sumType -> {
                var newSeenTypes = new HashSet<>(seenTypes);
                newSeenTypes.add(type);
                yield typesInfo.sumVariants(sumType).stream()
                    .anyMatch(variant -> isNonRecursivelyConstructable(variant.valueType(), newSeenTypes, typesInfo));
            }

            case TypeLevelValueType _ ->
                false;

            case TypeParam _ ->
                false;
        };
    }
}