package com.zaoo.lambda;

import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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 static final Logger log = LoggerFactory.getLogger(LambdaLocalServlet.class);
    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 {
            log.debug("invokeLambdaFunction:lambdaLocalPath={},lambdaFunction={}", lambdaFunction.getPath(), lambdaFunction.getPath());
            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);
        LambdaFunction ret = null;
        for (LambdaFunction lambdaFunction : lambdaFunctions) {
            if (path.startsWith(lambdaFunction.getPath())) {
                if (ret == null) {
                    ret = lambdaFunction;
                } else {
                    // There might be more than one lambdaFunction.path() matched. Pick the longest matched.
                    ret = lambdaFunction.getPath().length() > ret.getPath().length() ? lambdaFunction : ret;
                }
            }
        }
        return ret;
    }

    @SuppressWarnings("unchecked")
    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());
    }
}

@@ -1650,24 +1650,86 @@
 .toList());%0A
+        log.debug(%22initLambdaFunctions:%7B%7D%22, lambdaFunctions);%0A
     %7D%0A%0A    @
@@ -2104,53 +2104,9 @@
 h=%7B%7D
-,lambdaFunction=%7B%7D%22, lambdaFunction.getPath()
+%22
 , la
@@ -2775,24 +2775,128 @@
 unctions) %7B%0A
+            log.debug(%22findingLambdaFunction:path=%7B%7D,functionPath=%7B%7D%22, path, lambdaFunction.getPath());%0A
             

add more debug