package bench;

import java.io.File;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import megamek.client.bot.princess.Princess;
import megamek.common.Player;
import megamek.common.compute.Compute;
import megamek.common.enums.GamePhase;
import megamek.common.game.IGame;
import megamek.common.loaders.MekSummaryCache;
import megamek.common.scenario.Scenario;
import megamek.common.scenario.ScenarioLoader;
import megamek.server.Server;
import megamek.server.totalWarfare.TWGameManager;

/**
 * Density benchmark: how many concurrent MegaMek games fit in one JVM, and what
 * does game N+1 actually cost?
 *
 * <p>The design point that matters: games are started <em>staggered</em>, and each is
 * allowed to reach a steady state (past deployment, several rounds of real combat)
 * before the next one starts. A freshly constructed Server sitting in the lobby is
 * not representative - it has no turn history, no accumulated Report vectors, no
 * Princess path caches. Measuring those would flatter the result badly.
 *
 * <p>Usage:
 * <pre>
 *   MMBench --games 8 --scenario data/scenarios/... --settle-round 3 --max-rounds 30
 * </pre>
 */
public final class MMBench {

    private static final String LOCALHOST = "127.0.0.1";

    /** One running game: its server, its manager, and the bots playing it. */
    private static final class GameInstance {
        final int index;
        final int port;
        final Server server;
        final TWGameManager gameManager;
        final IGame game;
        final List<Princess> bots = new ArrayList<>();
        boolean finished;

        GameInstance(int index, int port, Server server, TWGameManager gameManager, IGame game) {
            this.index = index;
            this.port = port;
            this.server = server;
            this.gameManager = gameManager;
            this.game = game;
        }

        int round() {
            try {
                return game.getCurrentRound();
            } catch (Exception e) {
                return -1;
            }
        }

        GamePhase phase() {
            try {
                return game.getPhase();
            } catch (Exception e) {
                return null;
            }
        }

        boolean isOver() {
            GamePhase p = phase();
            return finished || p == GamePhase.VICTORY;
        }
    }

    public static void main(String[] args) throws Exception {
        Map<String, String> opts = parseArgs(args);

        int gameCount = Integer.parseInt(opts.getOrDefault("games", "4"));
        String scenarioPath = opts.getOrDefault("scenario",
              "data/scenarios/TrainingScenarios/1-FirstRun.mms");
        int basePort = Integer.parseInt(opts.getOrDefault("base-port", "8800"));
        int settleRound = Integer.parseInt(opts.getOrDefault("settle-round", "3"));
        int maxRounds = Integer.parseInt(opts.getOrDefault("max-rounds", "40"));
        int maxSeconds = Integer.parseInt(opts.getOrDefault("max-seconds", "900"));
        Path csv = Path.of(opts.getOrDefault("csv", "results/bench.csv"));

        System.out.printf("MMBench: games=%d scenario=%s settleRound=%d maxRounds=%d%n",
              gameCount, scenarioPath, settleRound, maxRounds);

        Sampler sampler = new Sampler(csv);
        sampler.checkpoint("jvm-start", 0, 0, 0);

        // Kick the RNG exactly as DedicatedServer does.
        Compute.d6();

        // Load the unit cache once. This is the big shared static allocation and the
        // whole reason multi-tenanting might pay off, so measure it in isolation.
        MekSummaryCache cache = MekSummaryCache.getInstance();
        while (!cache.isInitialized()) {
            Thread.sleep(250L);
        }
        System.out.printf("Unit cache initialized: %d units%n", cache.getAllMeks().length);
        sampler.checkpoint("unit-cache-loaded", 0, 0, 0);

        List<GameInstance> games = new ArrayList<>();
        long deadline = System.currentTimeMillis() + (maxSeconds * 1000L);

        for (int i = 0; i < gameCount; i++) {
            GameInstance instance = startGame(i, basePort + (i * 4), scenarioPath);
            games.add(instance);
            System.out.printf("-- started game %d on port %d with %d bots%n",
                  i, instance.port, instance.bots.size());

            // Let this game reach steady state before adding the next, so the
            // measured delta reflects a game that is actually playing.
            waitForSettle(instance, settleRound, sampler, games, deadline);
            sampler.checkpoint("settled-game-" + (i + 1), i + 1, liveCount(games), totalRounds(games));
        }

        System.out.println("All games started; running to completion...");
        int lastReported = -1;
        while (System.currentTimeMillis() < deadline && !allOver(games, maxRounds)) {
            Thread.sleep(2000L);
            int rounds = totalRounds(games);
            sampler.sample("running", games.size(), liveCount(games), rounds);
            int maxRound = games.stream().mapToInt(GameInstance::round).max().orElse(0);
            if (maxRound / 5 != lastReported) {
                lastReported = maxRound / 5;
                sampler.checkpoint("round-" + maxRound, games.size(), liveCount(games), rounds);
            }
        }

        sampler.checkpoint("final", games.size(), liveCount(games), totalRounds(games));

        System.out.println("\n=== per-game state ===");
        for (GameInstance g : games) {
            System.out.printf("game %d: round=%d phase=%s over=%s%n",
                  g.index, g.round(), g.phase(), g.isOver());
        }

        for (GameInstance g : games) {
            try {
                g.server.die();
            } catch (Exception ignored) {
                // best effort teardown
            }
        }

        System.out.println("Wrote " + csv.toAbsolutePath());
        System.exit(0);
    }

    /** Mirrors MegaMekGUI.scenario(), minus every GUI interaction. */
    private static GameInstance startGame(int index, int port, String scenarioPath) throws Exception {
        ScenarioLoader loader = new ScenarioLoader(new File(scenarioPath));
        Scenario scenario = loader.load();
        IGame game = scenario.createGame();

        TWGameManager gameManager = new TWGameManager();
        Server server = new Server("", port, gameManager);
        server.setGame(game);
        scenario.applyDamage(gameManager);
        gameManager.calculatePlayerInitialCounts();

        GameInstance instance = new GameInstance(index, port, server, gameManager, game);

        // Every player becomes a Princess: no human slot, so the game runs itself.
        for (Player player : game.getPlayersList()) {
            Princess bot = new Princess(player.getName(), LOCALHOST, port);
            bot.startPrecognition();
            bot.connect();
            instance.bots.add(bot);
        }

        // Bots chat in the lobby but never ready up on their own, so the game would
        // sit in LOUNGE forever. Push them out of it - but only once each bot has
        // actually received its entities.
        //
        // A fixed sleep here is a race: under CPU contention from already-running
        // games, a bot can still be waiting on its initial state when we ready it up,
        // and then it deploys with a null entity ("invalid deployment packet") and
        // the game wedges in DEPLOYMENT forever. Wait on observed state instead.
        long connectDeadline = System.currentTimeMillis() + 60_000L;
        for (Princess bot : instance.bots) {
            while (!bot.isConnected() && System.currentTimeMillis() < connectDeadline) {
                Thread.sleep(100L);
            }
        }
        for (Princess bot : instance.bots) {
            while (System.currentTimeMillis() < connectDeadline) {
                if (bot.getLocalPlayer() != null
                      && bot.getGame().getEntitiesOwnedBy(bot.getLocalPlayer()) > 0) {
                    break;
                }
                Thread.sleep(100L);
            }
            if (bot.getLocalPlayer() == null
                  || bot.getGame().getEntitiesOwnedBy(bot.getLocalPlayer()) == 0) {
                System.out.printf("!! bot %s never received entities before ready-up%n", bot.getName());
            }
        }
        for (Princess bot : instance.bots) {
            bot.sendDone(true);
        }

        return instance;
    }

    private static void waitForSettle(GameInstance instance, int settleRound, Sampler sampler,
          List<GameInstance> games, long deadline) throws InterruptedException {
        long settleDeadline = Math.min(deadline, System.currentTimeMillis() + 300_000L);
        while (System.currentTimeMillis() < settleDeadline) {
            if (instance.round() >= settleRound || instance.isOver()) {
                return;
            }
            Thread.sleep(1000L);
            sampler.sample("settling-" + instance.index, games.size(), liveCount(games), totalRounds(games));
        }
        System.out.printf("!! game %d did not reach round %d before timeout (round=%d phase=%s)%n",
              instance.index, settleRound, instance.round(), instance.phase());
    }

    private static int liveCount(List<GameInstance> games) {
        return (int) games.stream().filter(g -> !g.isOver()).count();
    }

    private static int totalRounds(List<GameInstance> games) {
        return games.stream().mapToInt(g -> Math.max(0, g.round())).sum();
    }

    private static boolean allOver(List<GameInstance> games, int maxRounds) {
        for (GameInstance g : games) {
            if (!g.isOver() && g.round() < maxRounds) {
                return false;
            }
        }
        return true;
    }

    private static Map<String, String> parseArgs(String[] args) {
        Map<String, String> map = new LinkedHashMap<>();
        for (int i = 0; i < args.length - 1; i++) {
            if (args[i].startsWith("--")) {
                map.put(args[i].substring(2), args[i + 1]);
            }
        }
        return map;
    }
}