package com.codeflash.profiler;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Configuration for the profiler agent, parsed from a JSON file.
 *
 * <p>The JSON is generated by Python ({@code JavaLineProfiler.generate_agent_config()}).
 * Uses a hand-rolled JSON parser to avoid external dependencies (keeps the agent JAR small).
 */
public final class ProfilerConfig {

    private String outputFile = "";
    private int warmupIterations = 10;
    private final Map<String, List<MethodTarget>> targets = new HashMap<>();
    private final Map<String, String> lineContents = new HashMap<>();
    private final Set<String> targetClassNames = new HashSet<>();

    public static class MethodTarget {
        public final String name;
        public final int startLine;
        public final int endLine;
        public final String sourceFile;

        public MethodTarget(String name, int startLine, int endLine, String sourceFile) {
            this.name = name;
            this.startLine = startLine;
            this.endLine = endLine;
            this.sourceFile = sourceFile;
        }
    }

    /**
     * Parse agent arguments and load the config file.
     *
     * <p>Expected format: {@code config=/path/to/config.json}
     */
    public static ProfilerConfig parse(String agentArgs) {
        ProfilerConfig config = new ProfilerConfig();
        if (agentArgs == null || agentArgs.isEmpty()) {
            return config;
        }

        String configPath = null;
        for (String part : agentArgs.split(",")) {
            String trimmed = part.trim();
            if (trimmed.startsWith("config=")) {
                configPath = trimmed.substring("config=".length());
            }
        }

        if (configPath == null) {
            System.err.println("[codeflash-profiler] No config= in agent args: " + agentArgs);
            return config;
        }

        try {
            String json = new String(Files.readAllBytes(Paths.get(configPath)), StandardCharsets.UTF_8);
            config.parseJson(json);
        } catch (IOException e) {
            System.err.println("[codeflash-profiler] Failed to read config: " + e.getMessage());
        }

        return config;
    }

    public String getOutputFile() {
        return outputFile;
    }

    public int getWarmupIterations() {
        return warmupIterations;
    }

    public Set<String> getTargetClasses() {
        return Collections.unmodifiableSet(targetClassNames);
    }

    public List<MethodTarget> getMethodsForClass(String internalClassName) {
        return targets.getOrDefault(internalClassName, Collections.emptyList());
    }

    public Map<String, String> getLineContents() {
        return Collections.unmodifiableMap(lineContents);
    }

    public int getExpectedLineCount() {
        int count = 0;
        for (List<MethodTarget> methods : targets.values()) {
            for (MethodTarget m : methods) {
                count += Math.max(m.endLine - m.startLine + 1, 1);
            }
        }
        return Math.max(count, 256);
    }

    /**
     * Check if a class should be instrumented. Uses JVM internal names (slash-separated).
     */
    public boolean shouldInstrumentClass(String internalClassName) {
        return targetClassNames.contains(internalClassName);
    }

    /**
     * Check if a specific method in a class should be instrumented.
     */
    public boolean shouldInstrumentMethod(String internalClassName, String methodName) {
        List<MethodTarget> methods = targets.get(internalClassName);
        if (methods == null) return false;
        for (MethodTarget m : methods) {
            if (m.name.equals(methodName)) {
                return true;
            }
        }
        return false;
    }

    /**
     * Resolve the absolute source file path for a given class and its source file attribute.
     */
    public String resolveSourceFile(String internalClassName) {
        List<MethodTarget> methods = targets.get(internalClassName);
        if (methods != null && !methods.isEmpty()) {
            return methods.get(0).sourceFile;
        }
        return internalClassName.replace('/', '.') + ".java";
    }

    // ---- Minimal JSON parser ----

    private void parseJson(String json) {
        json = json.trim();
        if (!json.startsWith("{") || !json.endsWith("}")) return;

        int[] pos = {1}; // mutable position cursor
        skipWhitespace(json, pos);

        while (pos[0] < json.length() - 1) {
            String key = readString(json, pos);
            skipWhitespace(json, pos);
            expect(json, pos, ':');
            skipWhitespace(json, pos);

            switch (key) {
                case "outputFile":
                    this.outputFile = readString(json, pos);
                    break;
                case "warmupIterations":
                    this.warmupIterations = readInt(json, pos);
                    break;
                case "targets":
                    parseTargets(json, pos);
                    break;
                case "lineContents":
                    parseLineContents(json, pos);
                    break;
                default:
                    skipValue(json, pos);
                    break;
            }

            skipWhitespace(json, pos);
            if (pos[0] < json.length() && json.charAt(pos[0]) == ',') {
                pos[0]++;
                skipWhitespace(json, pos);
            }
        }
    }

    private void parseTargets(String json, int[] pos) {
        expect(json, pos, '[');
        skipWhitespace(json, pos);

        while (pos[0] < json.length() && json.charAt(pos[0]) != ']') {
            parseTargetObject(json, pos);
            skipWhitespace(json, pos);
            if (pos[0] < json.length() && json.charAt(pos[0]) == ',') {
                pos[0]++;
                skipWhitespace(json, pos);
            }
        }
        pos[0]++; // skip ']'
    }

    private void parseTargetObject(String json, int[] pos) {
        expect(json, pos, '{');
        skipWhitespace(json, pos);

        String className = "";
        List<MethodTarget> methods = new ArrayList<>();

        while (pos[0] < json.length() && json.charAt(pos[0]) != '}') {
            String key = readString(json, pos);
            skipWhitespace(json, pos);
            expect(json, pos, ':');
            skipWhitespace(json, pos);

            switch (key) {
                case "className":
                    className = readString(json, pos);
                    break;
                case "methods":
                    methods = parseMethodsArray(json, pos);
                    break;
                default:
                    skipValue(json, pos);
                    break;
            }

            skipWhitespace(json, pos);
            if (pos[0] < json.length() && json.charAt(pos[0]) == ',') {
                pos[0]++;
                skipWhitespace(json, pos);
            }
        }
        pos[0]++; // skip '}'

        if (!className.isEmpty()) {
            targets.put(className, methods);
            targetClassNames.add(className);
        }
    }

    private List<MethodTarget> parseMethodsArray(String json, int[] pos) {
        List<MethodTarget> methods = new ArrayList<>();
        expect(json, pos, '[');
        skipWhitespace(json, pos);

        while (pos[0] < json.length() && json.charAt(pos[0]) != ']') {
            methods.add(parseMethodTarget(json, pos));
            skipWhitespace(json, pos);
            if (pos[0] < json.length() && json.charAt(pos[0]) == ',') {
                pos[0]++;
                skipWhitespace(json, pos);
            }
        }
        pos[0]++; // skip ']'
        return methods;
    }

    private MethodTarget parseMethodTarget(String json, int[] pos) {
        expect(json, pos, '{');
        skipWhitespace(json, pos);

        String name = "";
        int startLine = 0;
        int endLine = 0;
        String sourceFile = "";

        while (pos[0] < json.length() && json.charAt(pos[0]) != '}') {
            String key = readString(json, pos);
            skipWhitespace(json, pos);
            expect(json, pos, ':');
            skipWhitespace(json, pos);

            switch (key) {
                case "name":
                    name = readString(json, pos);
                    break;
                case "startLine":
                    startLine = readInt(json, pos);
                    break;
                case "endLine":
                    endLine = readInt(json, pos);
                    break;
                case "sourceFile":
                    sourceFile = readString(json, pos);
                    break;
                default:
                    skipValue(json, pos);
                    break;
            }

            skipWhitespace(json, pos);
            if (pos[0] < json.length() && json.charAt(pos[0]) == ',') {
                pos[0]++;
                skipWhitespace(json, pos);
            }
        }
        pos[0]++; // skip '}'

        return new MethodTarget(name, startLine, endLine, sourceFile);
    }

    private void parseLineContents(String json, int[] pos) {
        expect(json, pos, '{');
        skipWhitespace(json, pos);

        while (pos[0] < json.length() && json.charAt(pos[0]) != '}') {
            String key = readString(json, pos);
            skipWhitespace(json, pos);
            expect(json, pos, ':');
            skipWhitespace(json, pos);
            String value = readString(json, pos);
            lineContents.put(key, value);

            skipWhitespace(json, pos);
            if (pos[0] < json.length() && json.charAt(pos[0]) == ',') {
                pos[0]++;
                skipWhitespace(json, pos);
            }
        }
        pos[0]++; // skip '}'
    }

    private static String readString(String json, int[] pos) {
        if (pos[0] >= json.length() || json.charAt(pos[0]) != '"') return "";
        pos[0]++; // skip opening quote

        StringBuilder sb = new StringBuilder();
        while (pos[0] < json.length()) {
            char c = json.charAt(pos[0]);
            if (c == '\\' && pos[0] + 1 < json.length()) {
                pos[0]++;
                char escaped = json.charAt(pos[0]);
                switch (escaped) {
                    case '"': sb.append('"'); break;
                    case '\\': sb.append('\\'); break;
                    case '/': sb.append('/'); break;
                    case 'n': sb.append('\n'); break;
                    case 't': sb.append('\t'); break;
                    case 'r': sb.append('\r'); break;
                    default: sb.append('\\').append(escaped); break;
                }
            } else if (c == '"') {
                pos[0]++; // skip closing quote
                return sb.toString();
            } else {
                sb.append(c);
            }
            pos[0]++;
        }
        return sb.toString();
    }

    private static int readInt(String json, int[] pos) {
        int start = pos[0];
        boolean negative = false;
        if (pos[0] < json.length() && json.charAt(pos[0]) == '-') {
            negative = true;
            pos[0]++;
        }
        while (pos[0] < json.length() && Character.isDigit(json.charAt(pos[0]))) {
            pos[0]++;
        }
        String numStr = json.substring(start, pos[0]);
        try {
            return Integer.parseInt(numStr);
        } catch (NumberFormatException e) {
            return 0;
        }
    }

    private static void skipValue(String json, int[] pos) {
        if (pos[0] >= json.length()) return;
        char c = json.charAt(pos[0]);
        if (c == '"') {
            readString(json, pos);
        } else if (c == '{') {
            skipBraced(json, pos, '{', '}');
        } else if (c == '[') {
            skipBraced(json, pos, '[', ']');
        } else if (c == 'n' && json.startsWith("null", pos[0])) {
            pos[0] += 4;
        } else if (c == 't' && json.startsWith("true", pos[0])) {
            pos[0] += 4;
        } else if (c == 'f' && json.startsWith("false", pos[0])) {
            pos[0] += 5;
        } else {
            // number
            while (pos[0] < json.length() && "0123456789.eE+-".indexOf(json.charAt(pos[0])) >= 0) {
                pos[0]++;
            }
        }
    }

    private static void skipBraced(String json, int[] pos, char open, char close) {
        int depth = 0;
        boolean inString = false;
        while (pos[0] < json.length()) {
            char c = json.charAt(pos[0]);
            if (inString) {
                if (c == '\\') {
                    pos[0]++; // skip escaped char
                } else if (c == '"') {
                    inString = false;
                }
            } else {
                if (c == '"') inString = true;
                else if (c == open) depth++;
                else if (c == close) {
                    depth--;
                    if (depth == 0) {
                        pos[0]++;
                        return;
                    }
                }
            }
            pos[0]++;
        }
    }

    private static void skipWhitespace(String json, int[] pos) {
        while (pos[0] < json.length() && Character.isWhitespace(json.charAt(pos[0]))) {
            pos[0]++;
        }
    }

    private static void expect(String json, int[] pos, char expected) {
        if (pos[0] < json.length() && json.charAt(pos[0]) == expected) {
            pos[0]++;
        }
    }
}
