import megamek.common.CriticalSlot;
import megamek.common.battleValue.BVCalculator;
import megamek.common.equipment.AmmoMounted;
import megamek.common.equipment.Mounted;
import megamek.common.loaders.MekSummary;
import megamek.common.loaders.MekSummaryCache;
import megamek.common.units.Entity;
import megamek.common.units.Mek;

import java.io.BufferedWriter;
import java.io.FileOutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * Dump MegaMek's battle value for designs that have been shot at.
 *
 * <p>bv.jsonl is an oracle for designs as they leave the factory, which is the
 * only state a `.mtf` describes. It says nothing about what a Mek is worth once
 * its armour is gone, its gyro is cracked or its ammunition is spent - and that
 * is the case helm-bv exists to answer during a match, so it needs an oracle of
 * its own rather than a rule read off a page.
 *
 * <p>Each row is one design in one state: the damage applied, expressed the way
 * helm expresses it, and the battle value MegaMek computes for it. A fresh
 * Entity is loaded per scenario so no damage leaks into the next.
 *
 *   java -cp "MegaMek.jar:lib/*" DumpDamaged.java out.jsonl [design ...]
 *
 * Must run with MegaMek's install dir as the working directory.
 */
public final class DumpDamaged {

    /** Designs worth asking about: ordinary, ammunition-fed, and one of each awkward kind. */
    private static final String[] DEFAULT_DESIGNS = {
        "Atlas AS7-D",
        "Locust LCT-1V",
        "Warhammer WHM-6R",
        // MegaMek puts the reporting name first for a Clan design, so this is
        // the spelling its cache answers to.
        "Mad Cat (Timber Wolf) Prime",
        "Catapult CPLT-C1",
        "Hunchback HBK-4G",
        // One of each system whose loss changes a figure: jump jets, a hatchet
        // and the arm actuators that swing it, and MASC.
        "Griffin GRF-1N",
        "Hatchetman HCT-3F",
        "Celerity CLR-03-O",
        // A weapon that did not fit one location: the autocannon fills nine
        // slots of an arm and one of the torso beside it, so blowing the arm
        // off has to take the whole gun.
        "Victor VTR-9Ka",
        // An engine that does not run hot when it is hit, because it burns
        // fuel rather than fusing anything.
        "Pacifier CCU-36 SecurityMech",
    };

    public static void main(String[] args) throws Exception {
        if (args.length < 1) {
            System.err.println("usage: DumpDamaged <out.jsonl> [every:N | design ...]");
            System.exit(2);
        }
        // `every:N` samples the library rather than naming designs: every Nth
        // Mek by name, plus the hand-picked ones. The hand-picked list covers
        // the systems whose loss changes a figure; a sample covers the shapes
        // and construction the list does not - quads, tripods, superheavies,
        // Clan engines that reach into the side torsos, patchwork armour.
        int every = 0;
        String[] named = DEFAULT_DESIGNS;
        if (args.length > 1) {
            if (args[1].startsWith("every:")) {
                every = Integer.parseInt(args[1].substring("every:".length()));
            } else {
                named = java.util.Arrays.copyOfRange(args, 1, args.length);
            }
        }

        MekSummaryCache cache = MekSummaryCache.getInstance();
        long deadline = System.currentTimeMillis() + 600_000L;
        while (!cache.isInitialized()) {
            if (System.currentTimeMillis() > deadline) {
                System.err.println("units.cache did not initialise within 600s");
                System.exit(1);
            }
            Thread.sleep(250L);
        }

        Map<String, MekSummary> byName = new LinkedHashMap<>();
        for (MekSummary ms : cache.getAllMeks()) {
            byName.putIfAbsent(ms.getName(), ms);
        }

        java.util.List<String> wanted = new java.util.ArrayList<>(java.util.Arrays.asList(named));
        if (every > 0) {
            // Sorted so that the sample is the same one every time: the cache
            // iterates in whatever order it loaded, and an oracle that changes
            // under the test is the thing `logs/` already taught us to avoid.
            // Meks only. `getAllMeks` is named for the file format rather than
            // the unit type and hands back tanks and infantry too, whose
            // locations are not a Mek's - shooting a centre torso off one is
            // an index out of bounds rather than a scenario.
            java.util.List<String> all = new java.util.ArrayList<>();
            for (Map.Entry<String, MekSummary> e : byName.entrySet()) {
                if ("Mek".equalsIgnoreCase(e.getValue().getUnitType())) {
                    all.add(e.getKey());
                }
            }
            java.util.Collections.sort(all);
            for (int i = 0; i < all.size(); i += every) {
                if (!wanted.contains(all.get(i))) {
                    wanted.add(all.get(i));
                }
            }
        }

        int written = 0;
        try (PrintWriter out = new PrintWriter(new BufferedWriter(
                new OutputStreamWriter(new FileOutputStream(args[0]), StandardCharsets.UTF_8)))) {
            for (String name : wanted) {
                MekSummary ms = byName.get(name);
                if (ms == null) {
                    System.err.println("no such design: " + name);
                    continue;
                }
                for (String scenario : new String[] {
                        "fresh", "stripped", "torso-open", "structure-hurt", "arm-gone", "dry",
                        "hip-gone", "leg-actuators", "engine-hit", "gyro-hit", "heat-sinks-gone",
                        "jets-gone", "arm-actuators", "movement-gear-gone" }) {
                    Entity entity = ms.loadEntity();
                    if (entity == null) {
                        System.err.println("could not load: " + name);
                        break;
                    }
                    apply(entity, scenario);
                    BVCalculator calc = BVCalculator.getBVCalculator(entity);
                    int bv = calc.calculateBV(true, true);
                    out.println(row(name, scenario, bv, calc, entity));
                    written++;
                }
            }
        }
        System.err.println("wrote " + written + " rows");
    }

    /** Shoot the design in one specific way. */
    private static void apply(Entity entity, String scenario) {
        switch (scenario) {
            case "fresh" -> { }
            case "stripped" -> {
                for (int loc = 0; loc < entity.locations(); loc++) {
                    if (entity.getOArmor(loc) > 0) {
                        entity.setArmor(0, loc);
                    }
                    if (entity.hasRearArmor(loc) && entity.getOArmor(loc, true) > 0) {
                        entity.setArmor(0, loc, true);
                    }
                }
            }
            case "torso-open" -> {
                int ct = Mek.LOC_CENTER_TORSO;
                entity.setArmor(0, ct);
                if (entity.hasRearArmor(ct)) {
                    entity.setArmor(0, ct, true);
                }
            }
            case "structure-hurt" -> {
                // Half the frame gone in one torso, plate untouched, so the
                // structure term moves on its own.
                int lt = Mek.LOC_LEFT_TORSO;
                entity.setInternal(Math.max(1, entity.getOInternal(lt) / 2), lt);
            }
            case "arm-gone" -> destroyLocation(entity, Mek.LOC_RIGHT_ARM);
            // A leg's hip actuator shot out, and nothing else touched. In play
            // that halves what the Mek walks, and the question this asks is
            // whether the battle value follows.
            case "hip-gone" -> wreckSystem(entity, Mek.LOC_LEFT_LEG, Mek.ACTUATOR_HIP);
            // The three actuators below the hip, which each cost a point of
            // walking rather than halving it.
            case "leg-actuators" -> {
                wreckSystem(entity, Mek.LOC_LEFT_LEG, Mek.ACTUATOR_UPPER_LEG);
                wreckSystem(entity, Mek.LOC_LEFT_LEG, Mek.ACTUATOR_LOWER_LEG);
                wreckSystem(entity, Mek.LOC_LEFT_LEG, Mek.ACTUATOR_FOOT);
            }
            case "dry" -> {
                for (Mounted<?> m : entity.getEquipment()) {
                    if (m instanceof AmmoMounted ammo) {
                        ammo.setShotsLeft(0);
                    }
                }
            }
            // One engine hit, which is five points of heat a turn in play.
            case "engine-hit" -> wreckSystem(entity, Mek.LOC_CENTER_TORSO, Mek.SYSTEM_ENGINE);
            case "gyro-hit" -> wreckSystem(entity, Mek.LOC_CENTER_TORSO, Mek.SYSTEM_GYRO);
            // Cooling shot away, which is what makes a design run out of heat
            // sooner and count its weapons at half.
            case "heat-sinks-gone" -> wreckEquipment(entity, "Heat Sink");
            case "jets-gone" -> wreckEquipment(entity, "Jump Jet");
            // The two actuators a hatchet swings with, which the physical
            // attack values depend on.
            case "arm-actuators" -> {
                wreckSystem(entity, Mek.LOC_RIGHT_ARM, Mek.ACTUATOR_LOWER_ARM);
                wreckSystem(entity, Mek.LOC_RIGHT_ARM, Mek.ACTUATOR_HAND);
            }
            case "movement-gear-gone" -> {
                wreckEquipment(entity, "MASC");
                wreckEquipment(entity, "Supercharger");
                wreckEquipment(entity, "TSM");
            }
            default -> throw new IllegalArgumentException(scenario);
        }
    }

    /** Every mounted item whose name contains `what`, shot out where it sits. */
    private static void wreckEquipment(Entity entity, String what) {
        for (Mounted<?> m : entity.getEquipment()) {
            if (m.getType() == null || !m.getType().getInternalName().contains(what)) {
                continue;
            }
            int loc = m.getLocation();
            // Only what a critical hit could reach. A Mek's engine-integrated
            // heat sinks are mounted with no location and fill no critical
            // slot, so shooting them out describes damage that cannot happen
            // and that a .mul has no way to record - which would make this an
            // oracle for a state helm can never be asked about.
            if (loc < 0) {
                continue;
            }
            m.setDestroyed(true);
            m.setHit(true);
            for (int slot = 0; slot < entity.getNumberOfCriticalSlots(loc); slot++) {
                CriticalSlot cs = entity.getCritical(loc, slot);
                if (cs != null && cs.getType() == CriticalSlot.TYPE_EQUIPMENT
                        && cs.getMount() == m) {
                    cs.setDestroyed(true);
                    cs.setHit(true);
                }
            }
        }
    }

    /** One system critical slot shot out, leaving everything else alone. */
    private static void wreckSystem(Entity entity, int loc, int system) {
        for (int slot = 0; slot < entity.getNumberOfCriticalSlots(loc); slot++) {
            CriticalSlot cs = entity.getCritical(loc, slot);
            if (cs != null && cs.getType() == CriticalSlot.TYPE_SYSTEM
                    && cs.getIndex() == system) {
                cs.setDestroyed(true);
                cs.setHit(true);
            }
        }
    }

    /** Everything in one location shot out, the way a blown-off arm leaves it. */
    private static void destroyLocation(Entity entity, int loc) {
        for (int slot = 0; slot < entity.getNumberOfCriticalSlots(loc); slot++) {
            CriticalSlot cs = entity.getCritical(loc, slot);
            if (cs == null) {
                continue;
            }
            cs.setDestroyed(true);
            cs.setHit(true);
            if (cs.getMount() != null) {
                cs.getMount().setDestroyed(true);
                cs.getMount().setHit(true);
            }
        }
        entity.setArmor(0, loc);
        entity.setInternal(0, loc);
    }

    /** One row: the state as helm names it, and what MegaMek makes of it. */
    private static String row(String name, String scenario, int bv, BVCalculator calc,
            Entity entity) {
        StringBuilder sb = new StringBuilder(512);
        sb.append('{');
        str(sb, "name");
        sb.append(':');
        str(sb, name);
        sb.append(',');
        str(sb, "scenario");
        sb.append(':');
        str(sb, scenario);
        sb.append(",\"bv\":").append(bv);
        // Which half is wrong is the first question every disagreement asks.
        for (String field : new String[] { "defensiveValue", "offensiveValue" }) {
            try {
                java.lang.reflect.Field f = findField(calc.getClass(), field);
                if (f != null) {
                    f.setAccessible(true);
                    sb.append(",\"").append(field).append("\":").append(f.get(calc));
                }
            } catch (Throwable ignored) {
                // A MegaMek that renames these should lose the columns, not fail.
            }
        }

        sb.append(",\"armor\":{");
        boolean first = true;
        for (int loc = 0; loc < entity.locations(); loc++) {
            if (entity.getOArmor(loc) <= 0 && !entity.hasRearArmor(loc)) {
                continue;
            }
            if (!first) {
                sb.append(',');
            }
            first = false;
            str(sb, entity.getLocationAbbr(loc));
            sb.append(':').append(Math.max(0, entity.getArmor(loc)));
            if (entity.hasRearArmor(loc) && entity.getOArmor(loc, true) > 0) {
                sb.append(',');
                // The .mtf spells rear plate RTL, RTR and RTC.
                str(sb, "RT" + entity.getLocationAbbr(loc).charAt(0));
                sb.append(':').append(Math.max(0, entity.getArmor(loc, true)));
            }
        }
        sb.append("},\"structure\":{");
        first = true;
        for (int loc = 0; loc < entity.locations(); loc++) {
            if (!first) {
                sb.append(',');
            }
            first = false;
            str(sb, entity.getLocationAbbr(loc));
            sb.append(':').append(Math.max(0, entity.getInternal(loc)));
        }
        sb.append("},\"destroyed\":[");
        first = true;
        for (int loc = 0; loc < entity.locations(); loc++) {
            for (int slot = 0; slot < entity.getNumberOfCriticalSlots(loc); slot++) {
                CriticalSlot cs = entity.getCritical(loc, slot);
                if (cs == null || !cs.isDestroyed()) {
                    continue;
                }
                if (!first) {
                    sb.append(',');
                }
                first = false;
                sb.append('[');
                str(sb, entity.getLocationName(loc));
                sb.append(',').append(slot).append(']');
            }
        }
        sb.append("],\"emptyAmmo\":[");
        first = true;
        for (int loc = 0; loc < entity.locations(); loc++) {
            for (int slot = 0; slot < entity.getNumberOfCriticalSlots(loc); slot++) {
                CriticalSlot cs = entity.getCritical(loc, slot);
                if (cs == null || cs.getMount() == null) {
                    continue;
                }
                if (!(cs.getMount() instanceof AmmoMounted ammo) || ammo.getUsableShotsLeft() > 0) {
                    continue;
                }
                if (!first) {
                    sb.append(',');
                }
                first = false;
                sb.append('[');
                str(sb, entity.getLocationName(loc));
                sb.append(',').append(slot).append(']');
            }
        }
        sb.append("]}");
        return sb.toString();
    }

    private static java.lang.reflect.Field findField(Class<?> type, String name) {
        for (Class<?> c = type; c != null; c = c.getSuperclass()) {
            try {
                return c.getDeclaredField(name);
            } catch (NoSuchFieldException ignored) {
                // Keep walking up.
            }
        }
        return null;
    }

    private static void str(StringBuilder sb, String s) {
        sb.append('"');
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            switch (c) {
                case '"' -> sb.append("\\\"");
                case '\\' -> sb.append("\\\\");
                case '\n' -> sb.append("\\n");
                case '\r' -> sb.append("\\r");
                case '\t' -> sb.append("\\t");
                default -> {
                    if (c < 0x20) {
                        sb.append(String.format("\\u%04x", (int) c));
                    } else {
                        sb.append(c);
                    }
                }
            }
        }
        sb.append('"');
    }
}