package targetast;

import de.dhbwstuttgart.core.JavaTXCompiler;
import de.dhbwstuttgart.bytecode.Codegen;
import de.dhbwstuttgart.environment.DirectoryClassLoader;
import de.dhbwstuttgart.environment.IByteArrayClassLoader;
import de.dhbwstuttgart.parser.scope.JavaClassName;
import de.dhbwstuttgart.target.generate.ASTToTargetAST;
import de.dhbwstuttgart.target.tree.MethodParameter;
import de.dhbwstuttgart.target.tree.TargetClass;
import de.dhbwstuttgart.target.tree.TargetMethod;
import de.dhbwstuttgart.target.tree.TargetStructure;
import de.dhbwstuttgart.target.tree.expression.*;
import de.dhbwstuttgart.target.tree.type.TargetFunNType;
import de.dhbwstuttgart.target.tree.type.TargetRefType;
import de.dhbwstuttgart.target.tree.type.TargetType;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.*;

import org.objectweb.asm.Opcodes;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TestCodegen {
    static final Path outputPath = Path.of(System.getProperty("user.dir"), "/targetTest");

    private static void writeClassFile(String name, byte[] code) throws IOException {
        Files.createDirectories(outputPath);
        Files.write(outputPath.resolve(name + ".class"), code);
    }

    public static IByteArrayClassLoader createClassLoader() {
        return new DirectoryClassLoader(List.of(outputPath.toFile()), ClassLoader.getSystemClassLoader().getParent());
    }

    public static Path path = Path.of(System.getProperty("user.dir"), "resources/bytecode/javFiles/");

    public static Map<String, ? extends Class<?>> generateClassFiles(IByteArrayClassLoader classLoader, String... files) throws IOException, ClassNotFoundException {
        return generateClassFiles(classLoader, true, files);
    }

    public static Map<String, ? extends Class<?>> generateClassFiles(IByteArrayClassLoader classLoader, boolean inferTogether, String... files) throws IOException, ClassNotFoundException {
        Files.createDirectories(outputPath);
        var filenames = Arrays.stream(files).map(filename -> Path.of(path.toString(), filename).toFile()).toList();
        var compiler = new JavaTXCompiler(filenames, List.of(path.toFile(), outputPath.toFile()), outputPath.toFile(), inferTogether);
        compiler.generateBytecode();

        try(var newClassLoader = new DirectoryClassLoader(List.of(outputPath.toFile()), (ClassLoader)classLoader)) {
            var result = new HashMap<String, Class<?>>();
            for (var file : filenames) {
                var classes = compiler.sourceFiles.get(file).getClasses();

                result.putAll(classes.stream().map(cli -> {
                    try {
                        return newClassLoader.loadClass(cli.getClassName().toString());
                    } catch (ClassNotFoundException e) {
                        throw new RuntimeException(e);
                    }
                }).collect(Collectors.toMap(Class::getName, Function.identity())));
            }
            return result;
        }
    }

    public static Class<?> generateClass(TargetStructure clazz, IByteArrayClassLoader classLoader) throws IOException, ClassNotFoundException {
        Codegen codegen = new Codegen(clazz, new JavaTXCompiler(List.of()), null);
        var code = codegen.generate();
        writeClassFile(clazz.qualifiedName().getClassName(), code);
        return classLoader.loadClass(code);
    }

    public static Class<?> generateClass(TargetStructure clazz, IByteArrayClassLoader classLoader, ASTToTargetAST converter) throws IOException {
        Codegen codegen = new Codegen(clazz, converter.compiler, converter);
        var code = codegen.generate();
        writeClassFile(clazz.qualifiedName().getClassName(), code);
        return classLoader.loadClass(code);
    }

    public static Map<String, ? extends Class<?>> generateClassFiles(String filename, IByteArrayClassLoader classLoader) throws IOException, ClassNotFoundException {
        var file = Path.of(System.getProperty("user.dir"), "/resources/bytecode/javFiles/", filename).toFile();
        var compiler = new JavaTXCompiler(List.of(file), List.of(file.getParentFile()), outputPath.toFile(), true);
        var resultSet = compiler.typeInference(file);

        var sourceFile = compiler.sourceFiles.get(file);
        var converter = new ASTToTargetAST(compiler, resultSet, sourceFile, classLoader);
        var classes = compiler.sourceFiles.get(file).getClasses();

        var result = classes.stream().map(cli -> {
            try {
                return generateClass(converter.convert(cli), classLoader, converter);
            } catch (IOException exception) {
                throw new RuntimeException(exception);
            }
        }).collect(Collectors.toMap(Class::getName, Function.identity()));

        converter.generateFunNTypes();

        for (var entry : compiler.auxiliaries.entrySet()) {
            writeClassFile(entry.getKey(), entry.getValue());
        }

        return result;
    }

    @Test
    public void testEmptyClass() throws Exception {
        var clazz = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("Empty"));
        clazz.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "main", List.of(), null, new TargetBlock(List.of()));
        generateClass(clazz, createClassLoader()).getDeclaredMethod("main").invoke(null);
    }

    @Test
    public void testArithmetic() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("Arithmetic"));

        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "add", List.of(new MethodParameter(TargetType.Integer, "a"), new MethodParameter(TargetType.Integer, "b")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Add(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"), new TargetLocalVar(TargetType.Integer, "b"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "sub", List.of(new MethodParameter(TargetType.Integer, "a"), new MethodParameter(TargetType.Integer, "b")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Sub(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"), new TargetLocalVar(TargetType.Integer, "b"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "div", List.of(new MethodParameter(TargetType.Integer, "a"), new MethodParameter(TargetType.Integer, "b")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Div(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"), new TargetLocalVar(TargetType.Integer, "b"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "mul", List.of(new MethodParameter(TargetType.Integer, "a"), new MethodParameter(TargetType.Integer, "b")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Mul(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"), new TargetLocalVar(TargetType.Integer, "b"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "rem", List.of(new MethodParameter(TargetType.Integer, "a"), new MethodParameter(TargetType.Integer, "b")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Rem(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"), new TargetLocalVar(TargetType.Integer, "b"))))));

        var clazz = generateClass(targetClass, createClassLoader());
        assertEquals(clazz.getDeclaredMethod("add", Integer.class, Integer.class).invoke(null, 10, 10), 20);
        assertEquals(clazz.getDeclaredMethod("sub", Integer.class, Integer.class).invoke(null, 20, 10), 10);
        assertEquals(clazz.getDeclaredMethod("div", Integer.class, Integer.class).invoke(null, 20, 10), 2);
        assertEquals(clazz.getDeclaredMethod("mul", Integer.class, Integer.class).invoke(null, 20, 10), 200);
        assertEquals(clazz.getDeclaredMethod("rem", Integer.class, Integer.class).invoke(null, 10, 3), 1);
    }

    @Test
    public void testUnary() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("Unary"));

        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "not", List.of(new MethodParameter(TargetType.Integer, "a")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetUnaryOp.Not(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "neg", List.of(new MethodParameter(TargetType.Integer, "a")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetUnaryOp.Negate(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "add", List.of(new MethodParameter(TargetType.Integer, "a")), TargetType.Integer, new TargetBlock(List.of(new TargetReturn(new TargetUnaryOp.Add(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "a"))))));

        var clazz = generateClass(targetClass, createClassLoader());
        assertEquals(clazz.getDeclaredMethod("not", Integer.class).invoke(null, 10), -11);
        assertEquals(clazz.getDeclaredMethod("neg", Integer.class).invoke(null, 10), -10);
        assertEquals(clazz.getDeclaredMethod("add", Integer.class).invoke(null, 10), 10);

    }

    @Test
    public void testConditional() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("Conditional"));

        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "and", List.of(new MethodParameter(TargetType.Boolean, "a"), new MethodParameter(TargetType.Boolean, "b")), TargetType.Boolean, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.And(TargetType.Boolean, new TargetLocalVar(TargetType.Boolean, "a"), new TargetLocalVar(TargetType.Boolean, "b"))))));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "or", List.of(new MethodParameter(TargetType.Boolean, "a"), new MethodParameter(TargetType.Boolean, "b")), TargetType.Boolean, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Or(TargetType.Boolean, new TargetLocalVar(TargetType.Boolean, "a"), new TargetLocalVar(TargetType.Boolean, "b"))))));

        var clazz = generateClass(targetClass, createClassLoader());
        var and = clazz.getDeclaredMethod("and", Boolean.class, Boolean.class);
        var or = clazz.getDeclaredMethod("or", Boolean.class, Boolean.class);
        assertEquals(and.invoke(null, true, false), false);
        assertEquals(and.invoke(null, true, true), true);
        assertEquals(or.invoke(null, false, false), false);
        assertEquals(or.invoke(null, true, false), true);
    }

    // When adding two numbers and the return type is Long it needs to convert both values to Long
    @Test
    public void testArithmeticConvert() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("ArithmeticConvert"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "add", List.of(), TargetType.Long, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Add(TargetType.Long, new TargetLiteral.CharLiteral((char) 10), new TargetLiteral.LongLiteral((long) 20))))));
        var clazz = generateClass(targetClass, createClassLoader());
        assertEquals(clazz.getDeclaredMethod("add").invoke(null), (long) 30);
    }

    @Test
    public void testMethodCall() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("HelloWorld"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "helloWorld", List.of(), null, new TargetBlock(List.of(new TargetMethodCall(null, new TargetFieldVar(new TargetRefType("java.io.PrintStream"), new TargetRefType("java.lang.System"), true, new TargetClassName(new TargetRefType("java.lang.System")), "out"), List.of(new TargetLiteral.StringLiteral("Hello World!")), new TargetRefType("java.io.PrintStream"), "println", false, false, false))));

        var clazz = generateClass(targetClass, createClassLoader());
        clazz.getDeclaredMethod("helloWorld").invoke(null);
    }

    @Test
    public void testIfStatement() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("IfStmt"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "ifStmt", List.of(new MethodParameter(TargetType.Integer, "val")), TargetType.Integer, new TargetBlock(List.of(new TargetIf(new TargetBinaryOp.Equal(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "val"), new TargetLiteral.IntLiteral(10)), new TargetReturn(new TargetLiteral.IntLiteral(1)), new TargetIf(new TargetBinaryOp.Less(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "val"), new TargetLiteral.IntLiteral(5)), new TargetReturn(new TargetLiteral.IntLiteral(2)), new TargetReturn(new TargetLiteral.IntLiteral(3)))))));
        var clazz = generateClass(targetClass, createClassLoader());
        var ifStmt = clazz.getDeclaredMethod("ifStmt", Integer.class);
        assertEquals(ifStmt.invoke(null, 10), 1);
        assertEquals(ifStmt.invoke(null, 3), 2);
        assertEquals(ifStmt.invoke(null, 20), 3);
    }

    @Test
    public void testFor() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("For"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "forLoop", List.of(), TargetType.Integer, new TargetBlock(List.of(new TargetVarDecl(TargetType.Integer, "sum", new TargetLiteral.IntLiteral(0)), new TargetFor(List.of(new TargetVarDecl(TargetType.Integer, "i", new TargetLiteral.IntLiteral(0))), new TargetBinaryOp.Less(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetLiteral.IntLiteral(10)), List.of(new TargetAssign(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetBinaryOp.Add(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetLiteral.IntLiteral(1)))), new TargetBlock(List.of(new TargetAssign(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "sum"), new TargetBinaryOp.Add(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "sum"), new TargetLocalVar(TargetType.Integer, "i")))))), new TargetReturn(new TargetLocalVar(TargetType.Integer, "sum")))));
        var clazz = generateClass(targetClass, createClassLoader());
        assertEquals(clazz.getDeclaredMethod("forLoop").invoke(null), 45);
    }

    @Test
    public void testWhile() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("While"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "whileLoop", List.of(), TargetType.Integer, new TargetBlock(List.of(new TargetVarDecl(TargetType.Integer, "i", new TargetLiteral.IntLiteral(0)), new TargetWhile(new TargetBinaryOp.Less(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetLiteral.IntLiteral(10)), new TargetBlock(List.of(new TargetAssign(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetBinaryOp.Add(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetLiteral.IntLiteral(1)))))), new TargetReturn(new TargetLocalVar(TargetType.Integer, "i")))));
        var clazz = generateClass(targetClass, createClassLoader());
        assertEquals(clazz.getDeclaredMethod("whileLoop").invoke(null), 10);
    }

    @Test
    public void testClassicSwitch() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC , new JavaClassName("SwitchClassic"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "switchClassic", List.of(new MethodParameter(TargetType.Integer, "i")), TargetType.Integer, new TargetBlock(List.of(
                new TargetVarDecl(TargetType.Integer, "res", null),
                new TargetSwitch(new TargetLocalVar(TargetType.Integer, "i"), List.of(
                        new TargetSwitch.Case(List.of(new TargetLiteral.IntLiteral(10)), new TargetBlock(
                                List.of(new TargetAssign(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "res"), new TargetLiteral.IntLiteral(0)), new TargetBreak())
                        )),
                        new TargetSwitch.Case(List.of(new TargetLiteral.IntLiteral(15), new TargetLiteral.IntLiteral(20)), new TargetBlock(List.of())),
                        new TargetSwitch.Case(List.of(new TargetLiteral.IntLiteral(30)), new TargetBlock(
                                List.of(new TargetAssign(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "res"), new TargetLiteral.IntLiteral(1)), new TargetBreak())
                        ))
                ), new TargetSwitch.Case(new TargetBlock(
                        List.of(new TargetAssign(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "res"), new TargetLiteral.IntLiteral(2)), new TargetBreak())
                ), false)),
                new TargetReturn(new TargetLocalVar(TargetType.Integer, "res"))
        )));

        var clazz = generateClass(targetClass, createClassLoader());
        var m = clazz.getDeclaredMethod("switchClassic", Integer.class);
        assertEquals(m.invoke(null, 10), 0);
        assertEquals(m.invoke(null, 15), 1);
        assertEquals(m.invoke(null, 20), 1);
        assertEquals(m.invoke(null, 30), 1);
        assertEquals(m.invoke(null, 99), 2);
    }

    @Test
    public void testTypeSwitch() throws Exception {
        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("SwitchEnhanced"));
        targetClass.addMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "switchType", List.of(new MethodParameter(TargetType.Object, "obj")), TargetType.Integer, new TargetBlock(List.of(
                new TargetReturn(new TargetSwitch(new TargetLocalVar(TargetType.Object, "obj"), List.of(
                        new TargetSwitch.Case(List.of(new TargetTypePattern(TargetType.String, "aString")), new TargetBlock(
                                List.of(new TargetYield(new TargetLiteral.IntLiteral(0)))
                        ), false),
                        new TargetSwitch.Case(List.of(
                                new TargetGuard(new TargetTypePattern(TargetType.Integer, "i"), new TargetBinaryOp.Less(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "i"), new TargetLiteral.IntLiteral(10)))
                        ), new TargetBlock(
                                List.of(new TargetLiteral.IntLiteral(3))
                        ), true),
                        new TargetSwitch.Case(List.of(new TargetTypePattern(TargetType.Integer, "anInteger")), new TargetBlock(
                                List.of(new TargetLiteral.IntLiteral(1))
                        ), true)
                ), new TargetSwitch.Case(new TargetBlock(
                        List.of(new TargetLiteral.IntLiteral(2))
                ), true), TargetType.Integer)
                ))));
        var clazz = generateClass(targetClass, createClassLoader());
        var m = clazz.getDeclaredMethod("switchType", Object.class);
        assertEquals(m.invoke(null, "String"), 0);
        assertEquals(m.invoke(null,  10), 1);
        assertEquals(m.invoke(null, 'A'), 2);
        assertEquals(m.invoke(null,  5), 3);
    }

    @Test
    @Disabled("The lambda class is not generated because we don't call ASTToTargetAST")
    public void testLambda() throws Exception {
        var classLoader = createClassLoader();
        // var fun = classLoader.loadClass(Path.of(System.getProperty("user.dir"), "src/test/java/targetast/Fun1$$.class"));
        var interfaceType = TargetFunNType.fromParams(List.of(TargetType.Integer), 1);

        var targetClass = new TargetClass(Opcodes.ACC_PUBLIC, new JavaClassName("CGLambda"));
        targetClass.addConstructor(Opcodes.ACC_PUBLIC, List.of(), new TargetBlock(List.of(new TargetMethodCall(null, new TargetSuper(TargetType.Object), List.of(), TargetType.Object, "<init>", false, false, false))));
        var signature = new TargetMethod.Signature(Set.of(), List.of(new MethodParameter(TargetType.Integer, "num")), TargetType.Integer);
        targetClass.addMethod(Opcodes.ACC_PUBLIC, "lambda", List.of(), TargetType.Integer, new TargetBlock(List.of(new TargetVarDecl(interfaceType, "by2", new TargetLambdaExpression(interfaceType, List.of(), signature, new TargetBlock(List.of(new TargetReturn(new TargetBinaryOp.Mul(TargetType.Integer, new TargetLocalVar(TargetType.Integer, "num"), new TargetLiteral.IntLiteral(2))))))), new TargetReturn(new TargetCast(TargetType.Integer, new TargetMethodCall(TargetType.Object, TargetType.Object, List.of(TargetType.Object), new TargetLocalVar(interfaceType, "by2"), List.of(new TargetLiteral.IntLiteral(10)), interfaceType, "apply", false, true, false))))));
        var clazz = generateClass(targetClass, classLoader);
        var instance = clazz.getConstructor().newInstance();
        assertEquals(clazz.getDeclaredMethod("lambda").invoke(instance), 20);
    }
}
