package click.kamil.springstatemachineexporter.analysis.resolver;

import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.*;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.function.BiFunction;

/**
 * Compile-time literal folding for map lookups, array indexing, and narrow string transforms.
 */
final class LiteralExpressionSupport {

    private LiteralExpressionSupport() {
    }

    static String resolveArrayAccess(
            ArrayAccess access,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        Integer index = resolveLiteralInt(access.getIndex(), context, visited, resolver);
        if (index == null || index < 0) {
            return null;
        }
        List<String> elements = extractArrayElements(access.getArray(), context, visited, resolver);
        if (elements == null || index >= elements.size()) {
            return null;
        }
        return elements.get(index);
    }

    static String resolveMapGet(
            MethodInvocation getCall,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (getCall.arguments().isEmpty() || getCall.getExpression() == null) {
            return null;
        }
        String key = resolver.apply((Expression) getCall.arguments().get(0), visited);
        if (key == null) {
            return null;
        }
        Map<String, String> entries = extractMapEntries(getCall.getExpression(), context, visited, resolver);
        if (entries == null) {
            return null;
        }
        return entries.get(key);
    }

    static String resolveLiteralStringTransform(
            MethodInvocation transformCall,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (!transformCall.arguments().isEmpty() || transformCall.getExpression() == null) {
            return null;
        }
        String receiver = resolver.apply(transformCall.getExpression(), visited);
        if (receiver == null) {
            return null;
        }
        return switch (transformCall.getName().getIdentifier()) {
            case "toUpperCase" -> receiver.toUpperCase(Locale.ROOT);
            case "toLowerCase" -> receiver.toLowerCase(Locale.ROOT);
            case "trim" -> receiver.trim();
            default -> null;
        };
    }

    static Integer resolveLiteralInt(
            Expression expr,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (expr instanceof NumberLiteral nl) {
            try {
                return Integer.parseInt(nl.getToken().replace("_", ""));
            } catch (NumberFormatException e) {
                return null;
            }
        }
        if (expr instanceof ParenthesizedExpression pe) {
            return resolveLiteralInt(pe.getExpression(), context, visited, resolver);
        }
        if (expr instanceof InfixExpression infix) {
            Integer left = resolveLiteralInt(infix.getLeftOperand(), context, visited, resolver);
            Integer right = resolveLiteralInt(infix.getRightOperand(), context, visited, resolver);
            if (left == null || right == null) {
                return null;
            }
            InfixExpression.Operator op = infix.getOperator();
            if (op == InfixExpression.Operator.PLUS) {
                return left + right;
            }
            if (op == InfixExpression.Operator.MINUS) {
                return left - right;
            }
        }
        String resolved = resolver.apply(expr, visited);
        if (resolved != null) {
            try {
                return Integer.parseInt(resolved.replace("_", ""));
            } catch (NumberFormatException ignored) {
                return null;
            }
        }
        return null;
    }

    static Map<String, String> extractMapEntries(
            Expression expr,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (expr instanceof MethodInvocation mi) {
            Map<String, String> fromFactory = extractMapFactoryEntries(mi, context, visited, resolver);
            if (fromFactory != null) {
                return fromFactory;
            }
            return null;
        }
        String resolvedField = resolver.apply(expr, visited);
        if (resolvedField != null && resolvedField.startsWith("MAP:")) {
            return parseMapLiteral(resolvedField);
        }
        return null;
    }

    private static Map<String, String> extractMapFactoryEntries(
            MethodInvocation mi,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (!isMapFactory(mi.getExpression())) {
            return null;
        }
        String method = mi.getName().getIdentifier();
        if ("of".equals(method)) {
            return extractAlternatingPairs(mi.arguments(), context, visited, resolver);
        }
        if ("ofEntries".equals(method)) {
            Map<String, String> entries = new LinkedHashMap<>();
            for (Object argObj : mi.arguments()) {
                if (argObj instanceof Expression entryExpr) {
                    Map<String, String> single = extractMapEntryCall(entryExpr, context, visited, resolver);
                    if (single != null) {
                        entries.putAll(single);
                    }
                }
            }
            return entries.isEmpty() ? null : entries;
        }
        return null;
    }

    private static Map<String, String> extractMapEntryCall(
            Expression expr,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (!(expr instanceof MethodInvocation entryCall) || !"entry".equals(entryCall.getName().getIdentifier())) {
            return null;
        }
        if (!isMapFactory(entryCall.getExpression()) || entryCall.arguments().size() != 2) {
            return null;
        }
        String key = resolver.apply((Expression) entryCall.arguments().get(0), visited);
        String value = resolver.apply((Expression) entryCall.arguments().get(1), visited);
        if (key == null || value == null) {
            return null;
        }
        return Map.of(key, value);
    }

    private static Map<String, String> extractAlternatingPairs(
            List<?> arguments,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (arguments.size() < 2) {
            return null;
        }
        Map<String, String> entries = new LinkedHashMap<>();
        for (int i = 0; i + 1 < arguments.size(); i += 2) {
            String key = resolver.apply((Expression) arguments.get(i), visited);
            String value = resolver.apply((Expression) arguments.get(i + 1), visited);
            if (key == null || value == null) {
                return null;
            }
            entries.put(key, value);
        }
        return entries;
    }

    static String encodeMapLiteral(Map<String, String> entries) {
        StringBuilder sb = new StringBuilder("MAP:");
        boolean first = true;
        for (Map.Entry<String, String> entry : entries.entrySet()) {
            if (!first) {
                sb.append('|');
            }
            first = false;
            sb.append(escape(entry.getKey())).append('=').append(escape(entry.getValue()));
        }
        return sb.toString();
    }

    static Map<String, String> parseMapLiteral(String encoded) {
        if (encoded == null || !encoded.startsWith("MAP:")) {
            return null;
        }
        Map<String, String> entries = new LinkedHashMap<>();
        for (String pair : encoded.substring(4).split("\\|")) {
            int eq = pair.indexOf('=');
            if (eq <= 0) {
                continue;
            }
            entries.put(unescape(pair.substring(0, eq)), unescape(pair.substring(eq + 1)));
        }
        return entries.isEmpty() ? null : entries;
    }

    static List<String> mapLiteralValues(String encoded) {
        Map<String, String> entries = parseMapLiteral(encoded);
        if (entries == null) {
            return null;
        }
        return List.copyOf(entries.values());
    }

    private static String escape(String value) {
        return value.replace("\\", "\\\\").replace("|", "\\|").replace("=", "\\=");
    }

    private static String unescape(String value) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < value.length(); i++) {
            char c = value.charAt(i);
            if (c == '\\' && i + 1 < value.length()) {
                sb.append(value.charAt(++i));
            } else {
                sb.append(c);
            }
        }
        return sb.toString();
    }

    private static List<String> extractArrayElements(
            Expression arrayExpr,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        if (arrayExpr instanceof ArrayCreation creation && creation.getInitializer() != null) {
            return extractExpressionList(creation.getInitializer().expressions(), context, visited, resolver);
        }
        if (arrayExpr instanceof ArrayInitializer initializer) {
            return extractExpressionList(initializer.expressions(), context, visited, resolver);
        }
        String resolved = resolver.apply(arrayExpr, visited);
        if (resolved != null && resolved.startsWith("ARRAY:")) {
            return List.of(resolved.substring(6).split("\\|", -1));
        }
        return null;
    }

    static String encodeArrayLiteral(List<String> values) {
        return "ARRAY:" + String.join("|", values.stream().map(LiteralExpressionSupport::escape).toList());
    }

    private static List<String> extractExpressionList(
            List<?> expressions,
            CodebaseContext context,
            Set<String> visited,
            BiFunction<Expression, Set<String>, String> resolver) {
        List<String> values = new ArrayList<>();
        for (Object expressionObj : expressions) {
            String value = resolver.apply((Expression) expressionObj, visited);
            if (value == null) {
                return null;
            }
            values.add(value);
        }
        return values;
    }

    private static boolean isMapFactory(Expression expression) {
        if (expression == null) {
            return false;
        }
        if (expression instanceof SimpleName sn) {
            return "Map".equals(sn.getIdentifier());
        }
        if (expression instanceof QualifiedName qn) {
            return "Map".equals(qn.getName().getIdentifier());
        }
        return false;
    }
}
