package au.id.zancanaro.javacheck.junit;

import au.id.zancanaro.javacheck.Generator;
import au.id.zancanaro.javacheck.RoseTree;
import au.id.zancanaro.javacheck.ShrinkResult;
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.*;


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

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

    @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)) {
                continue;
            }
            Type type = field.getGenericType();
            if (!(type instanceof ParameterizedType)) {
                continue;
            }
            ParameterizedType ptype = (ParameterizedType) type;
            if (!(ptype.getRawType() instanceof Class)) {
                continue;
            }
            Class<?> c = (Class) ptype.getRawType();
            if (c != Generator.class) {
                continue;
            }
            boolean error = false;
            if (!Modifier.isStatic(field.getModifiers())) {
                errors.add(new Error("Generator field " + field.getName() + " must be static"));
                error = true;
            }
            if (!Modifier.isPublic(field.getModifiers())) {
                errors.add(new Error("Generator field " + field.getName() + " must be public"));
                error = true;
            }
            if (!error) {
                result.add(ptype.getActualTypeArguments()[0]);
            }
        }
        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 ptype = (ParameterizedType) type;
                if (!(ptype.getRawType() instanceof Class)) {
                    continue;
                }
                Class<?> c = (Class) ptype.getRawType();
                if (c != Generator.class) {
                    continue;
                }
                try {
                    Type target = ptype.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) {

                }
            }
        }
        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 {
                @SuppressWarnings("unchecked")
                Generator<?>[] generators = (Generator<?>[]) new Generator[method.getParameterCount()];
                int index = 0;
                for (Type type : method.getGenericParameterTypes()) {
                    generators[index++] = this.generators.get(type);
                }
                @SuppressWarnings("unchecked")
                Generator<Object[]> generator = Generator.tuple((Generator<Object>[]) generators).fmap(List::toArray);

                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);
                    RoseTree<Object[]> tree = generator.generate(random, size);
                    try {
                        runTest(tree.getValue());
                        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(RoseTree<Object[]> failed, Throwable originalEx) {
            ShrinkResult smallest = new ShrinkResult(failed.getValue(), originalEx);
            Iterator<RoseTree<Object[]>> trees = failed.getChildren();
            Set<List<Object>> seenArgs = new HashSet<>();
            while (trees.hasNext()) {
                RoseTree<Object[]> tree = trees.next();
                if (seenArgs.add(Arrays.asList(tree.getValue()))) {
                    try {
                        runTest(tree.getValue());
                    } catch (AssumptionViolatedException ex) {
                        // ignore, because it's not useful
                    } catch (Throwable ex) {
                        smallest = new ShrinkResult(tree.getValue(), ex);
                        Iterator<RoseTree<Object[]>> children = tree.getChildren();
                        if (children.hasNext()) {
                            trees = children;
                        } else {
                            break;
                        }
                    }
                }
            }
            return smallest;
        }

        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();
        }
    }

}