package com.zaoo.lambda;

import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
import org.springframework.util.AntPathMatcher;

import javax.servlet.ServletConfig;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.lang.reflect.InvocationTargetException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class LambdaLocalServlet extends HttpServlet {
    private final LambdaFunctionAnnotationScanner lambdaFunctionAnnotationScanner = new LambdaFunctionAnnotationScanner();
    private List<LambdaFunction> lambdaFunctions;

    @Override
    public void init(ServletConfig config) throws ServletException {
        super.init(config);

        String packages = config.getInitParameter("packages");
        if (packages == null) {
            throw new IllegalArgumentException("Unable to find 'init-param:packages' in servlet config");
        }

        initLambdaFunctions(Arrays.asList(packages.split(",")));
    }

    private void initLambdaFunctions(List<String> packageNames) {
        lambdaFunctions = packageNames.stream()
                .flatMap(packageName -> lambdaFunctionAnnotationScanner.listLambdaFunctions(packageName).stream())
                .collect(Collectors.toList());
    }

    @Override
    protected void service(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
        LambdaFunction lambdaFunction = getLambdaFunction(req);
        if (lambdaFunction == null) {
            super.service(req, resp);
            return;
        }

        try {
            invokeLambdaFunction(req, resp, lambdaFunction);
        } catch (InstantiationException | IllegalAccessException e) {
            throw new IllegalArgumentException("Unable to find handler class:" + lambdaFunction.getHandlerClass(), e);
        } catch (InvocationTargetException e) {
            throw new IllegalArgumentException("Unable to find handler method:" + lambdaFunction.getHandlerMethod(), e);
        }
    }

    private LambdaFunction getLambdaFunction(HttpServletRequest req) {
        String path = getRequestPath(req);
        for(LambdaFunction lambdaFunction : lambdaFunctions) {
            if(path.startsWith(lambdaFunction.getPath())) {
                return lambdaFunction;
            }
        }
        return null;
    }

    void invokeLambdaFunction(HttpServletRequest req,
                              HttpServletResponse resp,
                              LambdaFunction lambdaFunction) throws InstantiationException, IllegalAccessException, IOException, InvocationTargetException {
        Object handler = getLambdaFunctionObj(lambdaFunction);
        LambdaLocalContext context = new LambdaLocalContext();
        Object input = lambdaFunction.getDeserializer().serialize(req);

        if (handler instanceof RequestStreamHandler) {
            RequestStreamHandler streamHandler = (RequestStreamHandler) handler;
            byte[] buffer = (byte[]) input;
            try (ByteArrayInputStream in = new ByteArrayInputStream(buffer);
                 ByteArrayOutputStream out = new ByteArrayOutputStream()) {
                streamHandler.handleRequest(in, out, context);
                out.flush();
                lambdaFunction.getSerializer().deserialize(out.toByteArray(), resp);
            }
            return;
        }

        if (handler instanceof RequestHandler) {
            RequestHandler requestHandler = (RequestHandler) handler;
            Object output = requestHandler.handleRequest(input, context);
            lambdaFunction.getSerializer().deserialize(output, resp);
            return;
        }

        if (lambdaFunction.getHandlerMethod().getParameterCount() == 1) {
            Object output = lambdaFunction.getHandlerMethod().invoke(handler, input);
            lambdaFunction.getSerializer().deserialize(
                    lambdaFunction.getHandlerMethod().getReturnType().cast(output),
                    resp);
            return;
        }

        if (lambdaFunction.getHandlerMethod().getParameterCount() == 2) {
            Object output = lambdaFunction.getHandlerMethod().invoke(handler, input, context);
            lambdaFunction.getSerializer().deserialize(output, resp);
            return;
        }

        throw new IllegalStateException("Unable to find Lambda handler:" + lambdaFunction);
    }

    private Object getLambdaFunctionObj(LambdaFunction lambdaFunction) throws InstantiationException, IllegalAccessException {
        return lambdaFunction.getHandlerClass().newInstance();
    }

    private String getRequestPath(HttpServletRequest req) {
        return req.getRequestURI().substring(req.getContextPath().length());
    }
}

@@ -2376,24 +2376,59 @@
 tPath(req);%0A
+        LambdaFunction ret = null;%0A
         for(
@@ -2426,16 +2426,17 @@
      for
+ 
 (LambdaF
@@ -2493,16 +2493,17 @@
       if
+ 
 (path.st
@@ -2548,38 +2548,76 @@
                 
-return
+if (ret == null) %7B%0A                    ret =
  lambdaFunction;
@@ -2629,26 +2629,39 @@
         
-%7D%0A
+    %7D else %7B%0A  
         
 %7D%0A      
@@ -2652,26 +2652,26 @@
         
-%7D%0A
+  
         
 return n
@@ -2666,27 +2666,309 @@
     
-return null;%0A    %7D%0A
+// There might be more than one lambdaFunction.path() matched. Pick the longest matched.%0A                    ret = lambdaFunction.getPath().length() %3E ret.getPath().length() ? lambdaFunction : ret;%0A                %7D%0A            %7D%0A        %7D%0A        return ret;%0A    %7D%0A%0A    @SuppressWarnings(%22unchecked%22)
 %0A   

use greedy matches when matching @LambdaLocal.path()