import robocode.control.*;
import robocode.control.events.*;
import robocode.control.snapshot.*;

import java.io.*;
import java.util.*;

/**
 * Headless classic-Robocode battle runner that dumps per-turn robot state to JSONL.
 *
 * Usage: Capture <enemyNamePattern> <shooterNamePattern> <rounds> <outFile> <adversaryLabel> <note>
 *
 * The robot whose name matches enemyNamePattern is written as e* (the movement we want),
 * the other as s* (the shooter / bearing reference).
 *
 * Heading conversion (classic Robocode -> Tank Royale):
 *   classic heading: 0 = North (+y), positive = clockwise
 *   tank royale heading: 0 = East (+x), positive = counter-clockwise
 *   => headingTank = (90 - headingClassic) mod 360
 * Positions are identical (both origin bottom-left, x right, y up).
 */
public class Capture {

    public static class Row {
        int tick;            // global, continuous across rounds
        double ex, ey, eh, es, ee;
        double sx, sy, sh, ss, se;
    }

    static List<Row> rows = new ArrayList<>();
    static Map<Integer, String> idxToName = new HashMap<>();
    static int enemyIdx = -1, shooterIdx = -1;
    static int globalTick = 0;
    static Map<Integer, Integer> roundTicks = new LinkedHashMap<>();
    static List<String> errors = new ArrayList<>();

    static double tankHeading(double classicRad) {
        // IRobotSnapshot.getBodyHeading() returns RADIANS in the classic Robocode
        // convention (0 = north, positive = clockwise).
        double classicDeg = Math.toDegrees(classicRad);
        double t = (90.0 - classicDeg) % 360.0;
        if (t < 0) t += 360.0;
        return t;
    }

    static String f(double v) {
        return String.format(Locale.ROOT, "%.4f", v);
    }

    static boolean printedIndexMap = false;

    public static void main(String[] args) throws Exception {
        System.setProperty("java.awt.headless", "true");
        if (args.length < 6) {
            System.err.println("usage: Capture <enemyPattern> <shooterPattern> <rounds> <out> <label> <note>");
            System.exit(2);
        }
        final String enemyPattern = args[0];
        final String shooterPattern = args[1];
        final int rounds = Integer.parseInt(args[2]);
        final String outFile = args[3];
        final String label = args[4];
        final String note = args[5];

        File home = new File(System.getProperty("robocode.home", "/tmp/robocode/install"));
        System.out.println("[capture] home=" + home + " enemy=" + enemyPattern + " shooter=" + shooterPattern
                + " rounds=" + rounds);

        RobocodeEngine engine = new RobocodeEngine(home);
        engine.setVisible(false);
        engine.addBattleListener(new BattleAdaptor() {
            @Override
            public void onTurnEnded(TurnEndedEvent event) {
                try {
                    ITurnSnapshot snap = event.getTurnSnapshot();
                    IRobotSnapshot[] bots = snap.getRobots();
                    if (!printedIndexMap) {
                        printedIndexMap = true;
                        for (IRobotSnapshot b : bots) {
                            System.out.println("[capture] index=" + b.getRobotIndex() + " name=" + b.getName());
                        }
                    }
                    IRobotSnapshot en = null, sh = null;
                    for (IRobotSnapshot b : bots) {
                        String n = b.getName();
                        if (enemyIdx >= 0 && b.getRobotIndex() == enemyIdx) en = b;
                        if (shooterIdx >= 0 && b.getRobotIndex() == shooterIdx) sh = b;
                    }
                    if (en == null || sh == null) return;
                    if (en.getState() == null || !en.getState().isAlive()) return;
                    if (sh.getState() == null || !sh.getState().isAlive()) return;

                    Row r = new Row();
                    r.tick = globalTick++;
                    r.ex = en.getX(); r.ey = en.getY();
                    r.eh = tankHeading(en.getBodyHeading());
                    r.es = en.getVelocity(); r.ee = en.getEnergy();
                    r.sx = sh.getX(); r.sy = sh.getY();
                    r.sh = tankHeading(sh.getBodyHeading());
                    r.ss = sh.getVelocity(); r.se = sh.getEnergy();
                    rows.add(r);
                    roundTicks.merge(snap.getRound(), 1, Integer::sum);
                } catch (Throwable t) {
                    errors.add("onTurnEnded: " + t);
                }
            }

            @Override
            public void onRoundStarted(RoundStartedEvent e) {
                // nothing
            }

            @Override
            public void onBattleError(BattleErrorEvent e) {
                errors.add("battleError: " + e.getError());
            }

            @Override
            public void onRoundEnded(RoundEndedEvent e) {
                System.out.println("[capture] round " + e.getRound() + " ended, turns=" + e.getTurns());
            }
        });

        List<RobotSpecification> picked = new ArrayList<>();
        RobotSpecification enSpec = null, shSpec = null;
        String query = enemyPattern.equals(shooterPattern) ? enemyPattern : enemyPattern + "," + shooterPattern;
        RobotSpecification[] repo = engine.getLocalRepository(query);
        if (repo == null || repo.length < 1) {
            System.err.println("[capture] FAILED to resolve robots; got " + (repo == null ? 0 : repo.length));
            System.exit(3);
        }
        for (RobotSpecification rs : repo) {
            System.out.println("[capture] repo entry name=" + rs.getName() + " class=" + rs.getClassName());
            if (enSpec == null && rs.getClassName().equals(enemyPattern)) enSpec = rs;
            if (shSpec == null && rs.getClassName().equals(shooterPattern)) shSpec = rs;
        }
        if (shSpec == null && enemyPattern.equals(shooterPattern)) shSpec = enSpec;
        if (enSpec == null || shSpec == null) {
            System.err.println("[capture] could not match specs; enemy=" + enSpec + " shooter=" + shSpec);
            System.exit(3);
        }
        System.out.println("[capture] enemy spec = " + enSpec.getName());
        System.out.println("[capture] shooter spec = " + shSpec.getName());

        // Determine indices from the repository order passed to the battle.
        RobotSpecification[] chosen = new RobotSpecification[]{enSpec, shSpec};
        // The engine assigns robot index by order in the robots array (0-based).
        enemyIdx = 0;
        shooterIdx = 1;
        idxToName.put(0, enSpec.getName());
        idxToName.put(1, shSpec.getName());

        BattleSpecification spec = new BattleSpecification(
                rounds,                 // numRounds
                450L,                   // inactivity time (turns)
                0.1,                    // gun cooling rate
                new BattlefieldSpecification(800, 600),
                chosen);

        long t0 = System.currentTimeMillis();
        engine.runBattle(spec, true);
        engine.waitTillBattleOver();
        engine.close();
        long dt = System.currentTimeMillis() - t0;
        System.out.println("[capture] battle done in " + dt + " ms, rows=" + rows.size());

        // robust: verify index mapping by re-reading robot names once more is not possible; rely on order.

        // Write JSONL
        File of = new File(outFile);
        of.getParentFile().mkdirs();
        try (PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter(of)))) {
            StringBuilder meta = new StringBuilder();
            meta.append("{\"meta\":{");
            meta.append("\"adversary\":\"").append(label).append("\",");
            meta.append("\"source\":\"classic-robocode\",");
            meta.append("\"perfect_info\":true,");
            meta.append("\"arena\":{\"w\":800,\"h\":600},");
            meta.append("\"note\":\"").append(escape(note)).append("\"");
            meta.append("}}");
            out.println(meta);
            for (Row r : rows) {
                StringBuilder sb = new StringBuilder();
                sb.append("{\"tick\":").append(r.tick);
                sb.append(",\"ex\":").append(f(r.ex));
                sb.append(",\"ey\":").append(f(r.ey));
                sb.append(",\"eh\":").append(f(r.eh));
                sb.append(",\"es\":").append(f(r.es));
                sb.append(",\"ee\":").append(f(r.ee));
                sb.append(",\"sx\":").append(f(r.sx));
                sb.append(",\"sy\":").append(f(r.sy));
                sb.append(",\"sh\":").append(f(r.sh));
                sb.append(",\"ss\":").append(f(r.ss));
                sb.append(",\"se\":").append(f(r.se));
                sb.append("}");
                out.println(sb);
            }
        }
        System.out.println("[capture] wrote " + of.length() + " bytes to " + of.getAbsolutePath());

        // Sidecar with round boundaries (for statistics; not part of the shared fixture format).
        try (PrintWriter rs = new PrintWriter(new BufferedWriter(new FileWriter(outFile + ".rounds.json")))) {
            rs.print("{\"rounds\":[");
            int start = 0;
            boolean first = true;
            for (Map.Entry<Integer, Integer> e : roundTicks.entrySet()) {
                if (!first) rs.print(",");
                first = false;
                rs.print("{\"round\":" + e.getKey() + ",\"startTick\":" + start + ",\"count\":" + e.getValue() + "}");
                start += e.getValue();
            }
            rs.println("]}");
        }
        System.out.println("[capture] rounds ticks: " + roundTicks);
        if (!errors.isEmpty()) {
            System.out.println("[capture] errors (" + errors.size() + "):");
            for (String e : errors) System.out.println("   " + e);
        }
        // exit code 0 only if we captured something
        System.exit(rows.isEmpty() ? 4 : 0);
    }

    static String escape(String s) {
        return s.replace("\\", "\\\\").replace("\"", "\\\"");
    }
}