import dev.robocode.tankroyale.runner.*;
import dev.robocode.tankroyale.client.model.*;
import java.io.*;
import java.nio.file.*;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Training runner: PPO_Bot vs opponent for N rounds, all in one battle.
 *
 * All rounds run in a single battle so PPO_Bot's process (and its training
 * thread) survives across rounds — that is how checkpoints + round_counter
 * advance. Per-round battles killed the bot's training thread at round end.
 *
 * Writes one JSON line per round to the log file for game-level outcomes.
 * PPO_Bot itself writes training health fields (actorLoss, valueLoss, etc.)
 * to the same log file via PPOB_LOG_FILE — the shell wrapper stitches them.
 *
 * Usage (env vars):
 *   PPO_BOT_DIR      — path to the bot dir (any Tank Royale bot)
 *   SAMPLE_BOTS_DIR  — path to sample bots archive
 *   PPOB_LOG_FILE    — path to training_log.jsonl (appended)
 *   BOT_NAME         — bot name to match in round results (default: PPO_Bot)
 *   TRAINING_OPPONENT — opponent bot name (default: Target)
 *   TRAINING_ROUNDS  — number of rounds to run (CLI arg or env var)
 *
 * CLI: java RunTraining [opponent] [rounds]
 */
public class RunTraining {

    public static void main(String[] args) throws Exception {
        Logger.getLogger("dev.robocode.tankroyale").setLevel(Level.WARNING);

        String ppoBotDir  = requireEnv("PPO_BOT_DIR");
        String sampleBots = requireEnv("SAMPLE_BOTS_DIR");
        String logFile    = requireEnv("PPOB_LOG_FILE");

        String opponent = args.length > 0 ? args[0]
            : System.getenv().getOrDefault("TRAINING_OPPONENT", "Target");
        int totalRounds = args.length > 1 ? Integer.parseInt(args[1])
            : Integer.parseInt(System.getenv().getOrDefault("TRAINING_ROUNDS", "100"));
        String botName = System.getenv().getOrDefault("BOT_NAME", "PPO_Bot");

        System.out.printf("Training: %s vs %s for %d rounds%n", botName, opponent, totalRounds);
        System.out.printf("Log: %s%n", logFile);

        // Dead-bot guard: the runner keeps listing a crashed PPO_Bot in the
        // round results (name stays present, rank collapses), so `!found`
        // never fires. Instead watch the counter PPO_Bot writes at round
        // end (weights/round_counter.txt, BEFORE the runner's round event —
        // verified empirically). Aware that the counter write lags the
        // harness: sync training + checkpoint save take up to ~7s while
        // server rounds complete in 1-4s, so requiring only "no advance for
        // 2 harness rounds" false-positives on healthy bots (observed). A
        // truly dead process stays frozen for minutes, so abort only when
        // the counter has not advanced for BOTH >=10 harness rounds AND
        // >=10s wall time: exit(1) so run.sh's crash-restart loop resumes
        // from the counter. (Hoisted above the try so the end-of-battle
        // completeness check below can reuse them.)
        Path counterPath = Paths.get(ppoBotDir, "weights", "round_counter.txt");
        long[] lastCounter   = { readCounter(counterPath) };
        long    startCounter = lastCounter[0];
        int[]  frozenRounds  = { 0 };
        long[] lastAdvanceMs = { System.currentTimeMillis() };

        try (var runner = BattleRunner.create(b -> b.embeddedServer().suppressServerOutput())) {
            var setup = BattleSetup.classic(s -> s.setNumberOfRounds(totalRounds));
            var bots = List.of(
                BotEntry.of(ppoBotDir),
                BotEntry.of(sampleBots + "/" + opponent)
            );

            var owner = new Object();
            int[] prevTotal = { 0 };

            try (var handle = runner.startBattleAsync(setup, bots)) {

                handle.getOnRoundEnded().on(owner, event -> {
                    int round = event.getRoundNumber();
                    int ticks = event.getTurnNumber();
                    int totalScore = 0;
                    boolean win = false;
                    boolean found = false;
                    for (var r : event.getResults()) {
                        if (r.getName().equals(botName)) {
                            found = true;
                            totalScore = r.getTotalScore();
                            win = r.getRank() == 1;
                        }
                    }
                    int score = totalScore - prevTotal[0];
                    prevTotal[0] = totalScore;
                    // PPO_Bot's process died mid-battle: abort so run.sh's crash-restart
                    // loop resumes from round_counter instead of grinding dummy rounds.
                    // Frozen for a few harness rounds can be a healthy-but-lagging counter
                    // (sync training + checkpoint save vs fast server rounds), so also
                    // require 10s of wall time before declaring death.
                    long ctr = readCounter(counterPath);
                    if (ctr == lastCounter[0]) {
                        frozenRounds[0]++;
                        long frozenMs = System.currentTimeMillis() - lastAdvanceMs[0];
                        if (frozenRounds[0] >= 10 && frozenMs >= 10_000) {
                            System.err.printf("%s round_counter frozen at %d for %d "
                                + "harness rounds / %.0fs — process dead, aborting for restart%n",
                                botName, ctr, frozenRounds[0], frozenMs / 1000.0);
                            System.exit(1);
                        }
                    } else {
                        lastCounter[0] = ctr;
                        frozenRounds[0] = 0;
                        lastAdvanceMs[0] = System.currentTimeMillis();
                    }
                    if (!found) {
                        System.err.println(botName + " missing from round " + round
                            + " results — process died, aborting battle for restart");
                        System.exit(1);
                    }
                    // Append game-outcome JSON line
                    String line = String.format(
                        "{\"type\":\"game\",\"round\":%d,\"ticks\":%d,\"score\":%d,\"total_score\":%d,\"win\":%b,\"opponent\":\"%s\"}",
                        round, ticks, score, totalScore, win, opponent
                    );
                    try {
                        appendLine(logFile, line);
                    } catch (IOException e) {
                        System.err.println("Failed to append log line: " + e);
                    }
                    System.out.printf("Round %d/%d — ticks:%d score:%d win:%b%n",
                        round, totalRounds, ticks, score, win);
                });

                handle.awaitResults();

                // The final round's train lags the battle end: PPO_Bot's
                // onRoundEnded runs ppoUpdate + checkpoint save synchronously
                // on its own main thread, which can take longer than the
                // runner takes to deliver the last RoundEnded/GameEnded
                // events. So the counter right after awaitResults() is often
                // the stale pre-train value of the LAST round. Poll for it to
                // catch up before declaring the battle incomplete (the
                // mid-battle guard already catches truly frozen processes,
                // and a timeout still aborts for restart).
                long endCounter = readCounter(counterPath);
                long expectedEnd = startCounter + totalRounds;  // every round must have trained
                long waitUntil = System.currentTimeMillis() + 60_000;
                while (endCounter < expectedEnd && System.currentTimeMillis() < waitUntil) {
                    Thread.sleep(500);
                    endCounter = readCounter(counterPath);
                }
                if (endCounter < expectedEnd) {
                    System.err.printf("%s round_counter %d < expected %d (start+%d) at battle "
                        + "end (waited 60s) — %d rounds never trained (corpse?), aborting for "
                        + "restart%n",
                        botName, endCounter, expectedEnd, totalRounds, expectedEnd - endCounter);
                    System.exit(1);
                }
                System.out.printf("Counter check passed: %d == expected %d%n", endCounter, expectedEnd);
            }
        }
    }

    static void appendLine(String path, String line) throws IOException {
        try (var w = Files.newBufferedWriter(Paths.get(path),
                StandardOpenOption.CREATE,
                StandardOpenOption.APPEND)) {
            w.write(line);
            w.newLine();
        }
    }

    /** Read PPO_Bot's persisted round counter; -1 when absent/unreadable. */
    static long readCounter(Path p) {
        try {
            return Long.parseLong(Files.readString(p).trim());
        } catch (Exception e) {
            return -1;
        }
    }

    static String requireEnv(String name) {
        var v = System.getenv(name);
        if (v == null || v.isBlank()) {
            System.err.println("Error: " + name + " env var not set");
            System.exit(1);
        }
        return v;
    }
}