package au.id.zancanaro.javacheck.junit;

import au.id.zancanaro.javacheck.Generator;
import au.id.zancanaro.javacheck.ShrinkResult;
import au.id.zancanaro.javacheck.ShrinkTree;
import au.id.zancanaro.javacheck.annotations.DataSource;
import au.id.zancanaro.javacheck.annotations.Property;
import au.id.zancanaro.javacheck.annotations.Seed;
import org.junit.AssumptionViolatedException;
import org.junit.runners.BlockJUnit4ClassRunner;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.Statement;
import org.junit.runners.model.TestClass;

import java.lang.reflect.*;
import java.util.*;

@SuppressWarnings("WeakerAccess")
public class Properties extends BlockJUnit4ClassRunner {
    private final Map<Type, Generator<?>> generators = new HashMap<>();

    public Properties(Class<?> classObject) throws InitializationError {
        super(classObject);
    }

    @Override
    protected void collectInitializationErrors(List<Throwable> errors) {
        super.collectInitializationErrors(errors);
        Set<Type> generated = validateGeneratorFields(errors);
        validateTestMethodParameters(errors, generated);
    }

    private Set<Type> validateGeneratorFields(List<Throwable> errors) {
        Set<Type> result = new HashSet<>();
        Field[] fields = getTestClass().getJavaClass().getDeclaredFields();

        for (Field field : fields) {
            if (field.isAnnotationPresent(DataSource.class)) {
                boolean error = false;
                if (!Modifier.isStatic(field.getModifiers())) {
                    errors.add(new Error("@DataSource field " + field.getName() + " must be static"));
                    error = true;
                }
                if (!Modifier.isPublic(field.getModifiers())) {
                    errors.add(new Error("@DataSource field " + field.getName() + " must be public"));
                    error = true;
                }

                Type type = field.getGenericType();
                ParameterizedType parameterizedType;
                if (type instanceof ParameterizedType) {
                    parameterizedType = (ParameterizedType) type;
                    if (parameterizedType.getRawType() instanceof Class) {
                        Class<?> c = (Class) parameterizedType.getRawType();
                        if (c == Generator.class) {
                            if (!error) {
                                result.add(parameterizedType.getActualTypeArguments()[0]);
                            }
                        } else {
                            errors.add(new Error("@DataSource fields must be of type Generator<T>"));
                        }
                    } else {
                        errors.add(new Error("@DataSource fields must be of type Generator<T>"));
                    }
                } else {
                    errors.add(new Error("@DataSource fields must be of type Generator<T>"));
                }
            }
        }
        return result;
    }

    private void validateTestMethodParameters(List<Throwable> errors, Set<Type> generated) {
        for (FrameworkMethod each : computeTestMethods()) {
            for (Type type : each.getMethod().getGenericParameterTypes()) {
                if (!generated.contains(type)) {
                    errors.add(new Error("No @DataSource for type: " + type));
                    generated.add(type); // ignore future errors on this type
                }
            }
        }
    }

    private static final Map<Type, Type> rawTypes;

    static {
        Map<Type, Type> types = new HashMap<>();
        types.put(Double.class, Double.TYPE);
        types.put(Float.class, Float.TYPE);
        types.put(Long.class, Long.TYPE);
        types.put(Integer.class, Integer.TYPE);
        types.put(Short.class, Short.TYPE);
        types.put(Byte.class, Byte.TYPE);
        types.put(Character.class, Character.TYPE);
        types.put(Boolean.class, Boolean.TYPE);
        rawTypes = Collections.unmodifiableMap(types);
    }

    private Map<Type, Generator<?>> computeGenerators() {
        if (generators.isEmpty()) {
            Field[] fields = getTestClass().getJavaClass().getDeclaredFields();

            for (Field field : fields) {
                if (!field.isAnnotationPresent(DataSource.class)) {
                    continue;
                }
                Type type = field.getGenericType();
                if (!(type instanceof ParameterizedType)) {
                    continue;
                }
                ParameterizedType parameterizedType = (ParameterizedType) type;
                if (!(parameterizedType.getRawType() instanceof Class)) {
                    continue;
                }
                Class<?> c = (Class) parameterizedType.getRawType();
                if (c != Generator.class) {
                    continue;
                }
                try {
                    Type target = parameterizedType.getActualTypeArguments()[0];
                    @SuppressWarnings("unchecked")
                    Generator<Object> generator = (Generator<Object>) field.get(null);
                    generators.put(target, generator);
                    if (rawTypes.containsKey(target)) {
                        generators.put(rawTypes.get(target), generator);
                    }
                } catch (IllegalAccessException ex) {
                    throw new RuntimeException(ex);
                }
            }
        }
        return generators;
    }

    @Override
    protected void validateConstructor(List<Throwable> errors) {
        validateOnlyOneConstructor(errors);
    }

    @Override
    protected void validateTestMethods(List<Throwable> errors) {
        for (FrameworkMethod each : computeTestMethods()) {
            if (each.getAnnotation(Property.class) != null) {
                each.validatePublicVoid(false, errors);
                each.validateNoTypeParametersOnArgs(errors);
            } else {
                each.validatePublicVoidNoArg(false, errors);
            }
        }
    }

    @Override
    protected List<FrameworkMethod> computeTestMethods() {
        List<FrameworkMethod> testMethods = new ArrayList<>(super.computeTestMethods());
        List<FrameworkMethod> theoryMethods = getTestClass().getAnnotatedMethods(Property.class);
        testMethods.removeAll(theoryMethods);
        testMethods.addAll(theoryMethods);
        return testMethods;
    }

    @Override
    public Statement methodBlock(final FrameworkMethod method) {
        return new GenerativeTester(method, getTestClass(), computeGenerators());
    }

    public static class GenerativeTester extends Statement {
        private final FrameworkMethod testMethod;
        private final TestClass testClass;
        private final Map<Type, Generator<?>> generators;

        public GenerativeTester(FrameworkMethod testMethod, TestClass testClass, Map<Type, Generator<?>> generators) {
            this.testMethod = testMethod;
            this.testClass = testClass;
            this.generators = generators;
        }

        private long getSeed(Method method) {
            Seed seed = method.getAnnotation(Seed.class);
            if (seed == null) {
                return System.currentTimeMillis();
            } else {
                return seed.value();
            }
        }

        @Override
        public void evaluate() throws Throwable {
            Method method = testMethod.getMethod();
            if (method.getParameterCount() == 0) {
                runTest(new Object[0]);
            } else {
                Generator<?>[] generators = Arrays.stream(method.getGenericParameterTypes())
                        .map(this.generators::get)
                        .toArray(Generator[]::new);
                Generator<List<Object>> generator = Generator.tuple(generators);

                long seed = getSeed(method);
                Random random = new Random(seed);

                Property property = testMethod.getAnnotation(Property.class);
                int assumptionsViolated = 0;
                int maxSize = property.maxSize();
                int numTests = property.runs();
                for (int i = 0; i < numTests; ++i) {
                    int size = Math.min(i + 1, maxSize);
                    ShrinkTree<List<Object>> tree = generator.generate(random, size);
                    try {
                        runTest(tree.getValue().toArray());
                        assumptionsViolated = 0;
                    } catch (AssumptionViolatedException ex) {
                        numTests++;
                        if (assumptionsViolated++ == 50) {
                            throw new Error("Violated 50 assumptions in a row: failing test");
                        }
                    } catch (Throwable ex) {
//                        tree.print(new OutputStreamWriter(System.out), Arrays::toString);
                        throw new PropertyError(method.getName(), seed, shrink(tree, ex));
                    }
                }
            }
        }

        private ShrinkResult shrink(ShrinkTree<List<Object>> failed, Throwable originalEx) {
            // this array is a mutable container so the shutdown handler can see the new version
            ShrinkResult[] smallest = new ShrinkResult[]{
                    new ShrinkResult(failed.getValue(), originalEx)};

            Thread shutdownHandler = makeShutdownHandler(smallest, originalEx);
            Runtime.getRuntime().addShutdownHook(shutdownHandler);

            Iterator<ShrinkTree<List<Object>>> trees = failed.getChildren();
            Set<List<Object>> seenArgs = new HashSet<>();
            while (trees.hasNext()) {
                ShrinkTree<List<Object>> tree = trees.next();
                if (seenArgs.add(Arrays.asList(tree.getValue()))) {
                    try {
                        runTest(tree.getValue().toArray());
                    } catch (AssumptionViolatedException ex) {
                        // ignore, because it's not useful
                    } catch (Throwable ex) {
                        smallest[0] = new ShrinkResult(tree.getValue(), ex);
                        Iterator<ShrinkTree<List<Object>>> children = tree.getChildren();
                        if (children.hasNext()) {
                            trees = children;
                        } else {
                            break;
                        }
                    }
                }
            }

            Runtime.getRuntime().removeShutdownHook(shutdownHandler);
            return smallest[0];
        }

        private Thread makeShutdownHandler(ShrinkResult[] smallest, Throwable originalException) {
            return new Thread(() -> {
                System.err.println("Signal received while shrinking.\n" +
                        "Current best shrink is: " + smallest[0].args + "\n" +
                        "Shrinking exception: " + smallest[0].thrown + "\n" +
                        "Originally was: " + originalException);
            });
        }

        public void runTest(final Object[] args) throws Throwable {
            new BlockJUnit4ClassRunner(testClass.getJavaClass()) {
                @Override
                protected void collectInitializationErrors(
                        List<Throwable> errors) {
                    // do nothing
                }

                @Override
                public Statement methodBlock(FrameworkMethod method) {
                    return super.methodBlock(method);
                }

                @Override
                protected Statement methodInvoker(FrameworkMethod method, Object test) {
                    return new Statement() {
                        @Override
                        public void evaluate() throws Throwable {
                            method.invokeExplosively(test, args);
                        }
                    };
                }

                @Override
                public Object createTest() throws Exception {
                    return getTestClass().getOnlyConstructor().newInstance();
                }
            }.methodBlock(testMethod).evaluate();
        }
    }

}