package com.example.support;

import io.restassured.response.Response;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * A lightweight, owned load generator (L3git). Drives {@code users} concurrent virtual
 * users against a request - or a multi-step flow - for a duration, reusing the shared
 * {@link ApiContext}. Supports ramp-up (stagger worker starts), optional SLA thresholds
 * (p95 / error-rate), per-step stats, and an optional latency CSV for charting. Single JVM by
 * design (no distributed generators). Each worker uses {@link ApiContext#fire}, which does not
 * touch the shared {@code last}, so the run is thread-safe. Generated / shipped by L3git.
 */
public class LoadRunner {

    /** One request in a virtual user's flow. */
    public static final class Call {
        public final String method;
        public final String path;
        public final String body;
        public final String label;

        public Call(String method, String path, String body, String label) {
            this.method = method;
            this.path = path;
            this.body = body;
            this.label = (label == null || label.isEmpty()) ? (method.toUpperCase() + " " + path) : label;
        }
    }

    public static final class LoadOptions {
        public int users = 10;
        public long durationMs = 10000;
        public long rampMs = 0;
        public double maxP95Ms = -1;          // < 0 = no SLA on latency
        public double maxErrorRatePct = -1;   // < 0 = no SLA on errors
        public String csvPath = null;         // null/empty = no CSV
        public int agents = 1;                // ADVANCED: total load agents/machines (distributed)
        public int agent = 1;                 // this machine's 1-based index (usually via -Dload.agent)
        public final List<Call> flow = new ArrayList<>();
    }

    private static final class Sample {
        final long elapsedMs;
        final int latency;
        final int status;
        final boolean ok;
        final String label;

        Sample(long elapsedMs, int latency, int status, boolean ok, String label) {
            this.elapsedMs = elapsedMs;
            this.latency = latency;
            this.status = status;
            this.ok = ok;
            this.label = label;
        }
    }

    public static LoadResult run(ApiContext api, LoadOptions opt) {
        // ADVANCED - distributed sharding: run this node's slice of the total users. The SAME test runs
        // on N machines; each sets -Dload.agent=<i> (1..N) and -Dload.agents=<N> (or agents= in config),
        // so together they drive ~opt.users total. Each writes its own CSV (…​.agentN.csv) to aggregate.
        int agents = opt.agents, agent = opt.agent;
        try { String a = System.getProperty("load.agents"); if (a != null) agents = Integer.parseInt(a.trim()); } catch (Exception ignore) {}
        try { String a = System.getProperty("load.agent");  if (a != null) agent  = Integer.parseInt(a.trim()); } catch (Exception ignore) {}
        agents = Math.max(1, agents);
        agent = Math.max(1, Math.min(agent, agents));
        final int totalUsers = opt.users < 1 ? 1 : opt.users;
        final int workers = agents > 1 ? (int) Math.ceil((double) totalUsers / agents) : totalUsers;
        String csvPath = opt.csvPath;
        if (agents > 1) {
            System.out.printf("Distributed load: agent %d of %d - running %d of %d total users on this node.%n",
                    agent, agents, workers, totalUsers);
            if (csvPath != null && !csvPath.isEmpty()) {
                int dot = csvPath.lastIndexOf('.');
                csvPath = (dot > 0 ? csvPath.substring(0, dot) : csvPath) + ".agent" + agent + (dot > 0 ? csvPath.substring(dot) : "");
            }
        }
        final List<Call> flow = opt.flow.isEmpty()
                ? Collections.singletonList(new Call("GET", "/", null, null))
                : opt.flow;
        final long start = System.nanoTime();
        final long deadline = start + opt.durationMs * 1_000_000L;
        final List<Sample> samples = Collections.synchronizedList(new ArrayList<Sample>());

        ExecutorService pool = Executors.newFixedThreadPool(workers);
        List<Future<?>> futures = new ArrayList<>();
        for (int i = 0; i < workers; i++) {
            final long startDelayMs = opt.rampMs > 0 ? (long) ((double) i / workers * opt.rampMs) : 0;
            futures.add(pool.submit(() -> {
                if (startDelayMs > 0) {
                    try {
                        Thread.sleep(startDelayMs);
                    } catch (InterruptedException e) {
                        return;
                    }
                }
                while (System.nanoTime() < deadline) {
                    for (Call c : flow) {
                        if (System.nanoTime() >= deadline) {
                            break;
                        }
                        long t0 = System.nanoTime();
                        boolean ok;
                        int sc;
                        try {
                            Response r = api.fire(c.method, c.path, c.body);
                            sc = r.getStatusCode();
                            ok = sc >= 200 && sc < 400;
                        } catch (Exception e) {
                            sc = 0;
                            ok = false;
                        }
                        long lat = (System.nanoTime() - t0) / 1_000_000L;
                        long elapsed = (System.nanoTime() - start) / 1_000_000L;
                        samples.add(new Sample(elapsed, (int) lat, sc, ok, c.label));
                    }
                }
            }));
        }
        for (Future<?> f : futures) {
            try {
                f.get();
            } catch (Exception e) {
                // worker failures already recorded as errors inside the loop
            }
        }
        pool.shutdownNow();
        long elapsedMs = (System.nanoTime() - start) / 1_000_000L;

        // --- aggregate ---
        List<Long> latencies = new ArrayList<>(samples.size());
        int errors = 0;
        Map<String, List<Long>> byLabel = new LinkedHashMap<>();
        Map<String, int[]> labelCounts = new LinkedHashMap<>();   // [count, errors]
        for (Sample s : samples) {
            latencies.add((long) s.latency);
            if (!s.ok) {
                errors++;
            }
            byLabel.computeIfAbsent(s.label, k -> new ArrayList<>()).add((long) s.latency);
            int[] c = labelCounts.computeIfAbsent(s.label, k -> new int[2]);
            c[0]++;
            if (!s.ok) {
                c[1]++;
            }
        }
        List<LoadResult.StepStat> steps = new ArrayList<>();
        for (Map.Entry<String, List<Long>> e : byLabel.entrySet()) {
            List<Long> ls = new ArrayList<>(e.getValue());
            Collections.sort(ls);
            long sum = 0;
            for (long l : ls) {
                sum += l;
            }
            double avg = ls.isEmpty() ? 0 : (double) sum / ls.size();
            int[] c = labelCounts.get(e.getKey());
            steps.add(new LoadResult.StepStat(e.getKey(), c[0], c[1],
                    LoadResult.pct(ls, 50), LoadResult.pct(ls, 95), avg));
        }

        // --- SLA verdict (only when a threshold was set) ---
        boolean slaChecked = opt.maxP95Ms >= 0 || opt.maxErrorRatePct >= 0;
        boolean slaPassed = true;
        StringBuilder sla = new StringBuilder();
        List<Long> sorted = new ArrayList<>(latencies);
        Collections.sort(sorted);
        long p95all = LoadResult.pct(sorted, 95);
        double errRate = samples.isEmpty() ? 0 : (errors * 100.0 / samples.size());
        if (opt.maxP95Ms >= 0) {
            boolean ok = p95all <= opt.maxP95Ms;
            slaPassed = slaPassed && ok;
            sla.append(String.format("p95 %d ms %s %.0f ms", p95all, ok ? "<=" : ">", opt.maxP95Ms));
        }
        if (opt.maxErrorRatePct >= 0) {
            boolean ok = errRate <= opt.maxErrorRatePct;
            slaPassed = slaPassed && ok;
            if (sla.length() > 0) {
                sla.append("; ");
            }
            sla.append(String.format("errors %.1f%% %s %.1f%%", errRate, ok ? "<=" : ">", opt.maxErrorRatePct));
        }

        if (csvPath != null && !csvPath.isEmpty()) {
            writeCsv(csvPath, samples);
        }
        return new LoadResult(samples.size(), errors, elapsedMs, latencies, steps,
                slaChecked, slaPassed, sla.toString());
    }

    private static void writeCsv(String path, List<Sample> samples) {
        List<Sample> sorted = new ArrayList<>(samples);
        sorted.sort((a, b) -> Long.compare(a.elapsedMs, b.elapsedMs));
        File f = new File(path);
        if (f.getParentFile() != null) {
            f.getParentFile().mkdirs();
        }
        try (PrintWriter w = new PrintWriter(new FileWriter(f))) {
            w.println("elapsed_ms,latency_ms,status,ok,label");
            for (Sample s : sorted) {
                w.printf("%d,%d,%d,%s,%s%n", s.elapsedMs, s.latency, s.status, s.ok, csv(s.label));
            }
        } catch (IOException e) {
            System.out.println("Load CSV write failed: " + e.getMessage());
        }
    }

    private static String csv(String v) {
        return v.contains(",") || v.contains("\"") ? "\"" + v.replace("\"", "\"\"") + "\"" : v;
    }

    /**
     * Parse a plain-text load config (a Cucumber doc string). One directive per line:
     * <pre>
     *   users = 50
     *   seconds = 30
     *   ramp = 10                 (optional; 0 = all users start at once)
     *   p95 = 800                 (optional SLA: fail if p95 latency &gt; 800 ms)
     *   errors = 1                (optional SLA: fail if error rate &gt; 1%)
     *   csv = target/load.csv     (optional: write per-request samples for charting)
     *   call GET /catalog         (one or more; body is the rest of the line, single-line)
     *   call POST /cart {"item":1}
     * </pre>
     * Unknown / malformed lines are ignored (best-effort).
     */
    public static LoadOptions parse(String config) {
        LoadOptions o = new LoadOptions();
        if (config == null) {
            return o;
        }
        for (String raw : config.split("\\r?\\n")) {
            String line = raw.trim();
            if (line.isEmpty() || line.startsWith("#")) {
                continue;
            }
            if (line.toLowerCase().startsWith("call ")) {
                String[] parts = line.substring(5).trim().split("\\s+", 3);
                if (parts.length >= 2) {
                    String body = parts.length >= 3 ? parts[2] : null;
                    o.flow.add(new Call(parts[0], parts[1], body, null));
                }
                continue;
            }
            int eq = line.indexOf('=');
            if (eq < 0) {
                continue;
            }
            String key = line.substring(0, eq).trim().toLowerCase();
            String val = line.substring(eq + 1).trim();
            try {
                switch (key) {
                    case "users": o.users = Integer.parseInt(val); break;
                    case "seconds": o.durationMs = (long) (Double.parseDouble(val) * 1000); break;
                    case "ms": o.durationMs = Long.parseLong(val); break;
                    case "ramp": o.rampMs = (long) (Double.parseDouble(val) * 1000); break;
                    case "p95": o.maxP95Ms = Double.parseDouble(val); break;
                    case "errors": o.maxErrorRatePct = Double.parseDouble(val); break;
                    case "csv": o.csvPath = val; break;
                    case "agents": o.agents = Integer.parseInt(val); break;
                    default: break;
                }
            } catch (NumberFormatException e) {
                // ignore a malformed number
            }
        }
        return o;
    }

    // --- back-compat single-request overloads ---
    public static LoadResult run(ApiContext api, String method, String path, String body, int users, long durationMs) {
        LoadOptions o = new LoadOptions();
        o.users = users;
        o.durationMs = durationMs;
        o.flow.add(new Call(method, path, body, null));
        return run(api, o);
    }

    public static LoadResult run(ApiContext api, String method, String path, String body, int users, int seconds) {
        return run(api, method, path, body, users, seconds * 1000L);
    }
}
