import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * De agent-lus — mijlpaal M8: het model krijgt een opdracht en gereedschap,
 * en werkt zelfstandig tot de deterministische controles tevreden zijn.
 *
 * <pre>
 *   java --add-modules jdk.incubator.vector src/Agent.java model.gguf \
 *        --taak "Bouw pagina.html: een pagina over koffie met een kop en drie soorten in een lijst." \
 *        --eis "&lt;h1&gt;" --eis "&lt;ul&gt;"
 * </pre>
 *
 * <p>Het ontwerp draait om &eacute;&eacute;n gedachte: <b>een klein model is wiebelig,
 * dus de zekerheid moet uit het harnas komen, niet uit het model.</b>
 * <ol>
 *   <li><b>Weinig gereedschap.</b> Twee stuks: {@code write_file} en {@code finish}.
 *       Elk extra gereedschap is een extra kans op de verkeerde keuze.</li>
 *   <li><b>Controle na elke stap.</b> Na elke schrijfactie draait het harnas
 *       {@link HtmlControles} en krijgt het model de uitkomst te zien &mdash; het hoeft
 *       er niet om te vragen, en kan er niet omheen.</li>
 *   <li><b>Een poortwachter op het einde.</b> {@code finish} wordt geweigerd zolang
 *       er problemen openstaan. Het model bepaalt de weg, het harnas bepaalt
 *       wanneer het af is.</li>
 * </ol>
 *
 * <p>De gereedschapsnamen en -beschrijvingen zijn Engels: het model heeft tijdens
 * de training vrijwel uitsluitend Engelse schema's gezien, en bij een half miljard
 * parameters telt elke vertrouwde vorm. De opdracht zelf mag gewoon Nederlands zijn.
 */
public final class Agent {

    static final String GEBRUIK = "gebruik: java --add-modules jdk.incubator.vector src/Agent.java "
            + "<model.gguf> --taak \"...\" [--map agentwerk] [--eis \"...\"]... [--stappen 12] [--ctx N]";

    // De twee gereedschapsschema's, in het formaat dat het sjabloon voorschrijft.
    static final String SCHEMA_WRITE = """
            {"type": "function", "function": {"name": "write_file", "description": "Write a complete file. Overwrites the file if it exists. The result reports any problems found by automatic checks.", "parameters": {"type": "object", "properties": {"path": {"type": "string", "description": "File name, for example page.html"}, "content": {"type": "string", "description": "The complete file content"}}, "required": ["path", "content"]}}}""";
    static final String SCHEMA_FINISH = """
            {"type": "function", "function": {"name": "finish", "description": "Finish the task. Only call this after write_file reported: no problems found.", "parameters": {"type": "object", "properties": {"summary": {"type": "string", "description": "One sentence describing the result"}}, "required": ["summary"]}}}""";

    static final String SYSTEEM = "You are a careful web developer. Complete the user's task by calling "
            + "the provided tools. Always respond with exactly one <tool_call> block and nothing else.";

    // Twee steuntjes voor een klein model, elk afzonderlijk uitschakelbaar voor de meting:
    // (1) een voorgekauwd voorbeeldgesprek dat het hele protocol één keer voordoet;
    // (2) een voorzet: het harnas begint elk antwoord alvast met "<tool_call>\n",
    //     zodat "sorry, dat kan ik niet" geen bereikbare uitweg meer is.
    static final String VOORZET = "<tool_call>\n";
    static final String VOORBEELD_TAAK = "Example task: create hello.html, a page with a greeting.";
    static final String VOORBEELD_CALL1 = "<tool_call>\n{\"name\": \"write_file\", \"arguments\": "
            + "{\"path\": \"hello.html\", \"content\": \"<!DOCTYPE html>\\n<html lang=\\\"en\\\">\\n"
            + "<head><title>Hello</title></head>\\n<body><h1>Hello!</h1></body>\\n</html>\"}}\n</tool_call>";
    static final String VOORBEELD_CALL2 = "<tool_call>\n{\"name\": \"finish\", \"arguments\": "
            + "{\"summary\": \"hello.html created and checked.\"}}\n</tool_call>";

    // MAX_ANTWOORD begrenst het antwoord op 1200 tokens (~5 KB tekst); MAX_TEKENS is
    // het vangnet daarboven, dus bewust in dezelfde orde van grootte
    static final int MAX_BESTANDEN = 12, MAX_TEKENS = 20_000, MAX_ANTWOORD = 1200;

    public static void main(String[] args) throws Exception {
        if (args.length < 1) { System.err.println(GEBRUIK); System.exit(2); }
        Path model = Path.of(args[0]);
        String taak = null;
        Path map = Path.of("agentwerk");
        List<String> eisen = new ArrayList<>();
        int stappen = 12, ctx = 4096;
        boolean metVoorbeeld = true, metVoorzet = true;
        for (int i = 1; i < args.length; i++) {
            switch (args[i]) {
                case "--taak"    -> taak = eis(args, ++i, "--taak");
                case "--map"     -> map = Path.of(eis(args, ++i, "--map"));
                case "--eis"     -> {
                    String e = eis(args, ++i, "--eis");
                    if (e.contains("<!--") || e.contains("-->")) {
                        // commentaar wordt vóór de controle weggeknipt; zo'n eis kan nooit slagen
                        System.err.println("--eis mag geen commentaar-syntaxis bevatten: " + e);
                        System.exit(2);
                    }
                    eisen.add(e);
                }
                case "--stappen" -> stappen = Integer.parseInt(eis(args, ++i, "--stappen"));
                case "--ctx"     -> ctx = Integer.parseInt(eis(args, ++i, "--ctx"));
                case "--zonder-voorbeeld" -> metVoorbeeld = false;
                case "--zonder-voorzet"   -> metVoorzet = false;
                default -> { System.err.println("onbekende optie: " + args[i]); System.exit(2); }
            }
        }
        if (taak == null) { System.err.println("--taak ontbreekt.\n" + GEBRUIK); System.exit(2); }
        Files.createDirectories(map);

        try (Gguf g = Gguf.open(model)) {
            if (!g.architecture().equals("qwen2")) {
                // qwen3.5/3.8 bieden gereedschap aan in een XML-dialect; dit harnas
                // spreekt bewust nog het qwen2-JSON-dialect (zie de LEESMIJ bij M9)
                System.err.println("de agent ondersteunt voorlopig alleen qwen2-modellen; de "
                        + "qwen3.5/3.8-familie gebruikt een ander gereedschapsdialect.");
                System.exit(2);
            }
            Tokenizer tk = new Tokenizer(g);
            ChatTemplate sjabloon = ChatTemplate.van(g);
            Motor m = Motor.open(g, ctx);
            Sampler sampler = Sampler.greedy();          // herhaalbaarheid boven creativiteit

            System.out.printf("agent  : %s  |  context %d, max %d stappen%n",
                    g.has("general.name") ? g.getString("general.name") : model.getFileName(), m.ctx(), stappen);
            System.out.println("map    : " + map + "  |  eisen: " + (eisen.isEmpty() ? "(geen)" : eisen));
            System.out.println("taak   : " + taak);

            StringBuilder open = new StringBuilder(
                    sjabloon.openingMetTools(null, List.of(SCHEMA_WRITE, SCHEMA_FINISH)));
            if (metVoorbeeld) {
                open.append(sjabloon.blok("user", VOORBEELD_TAAK))
                    .append("<|im_start|>assistant\n").append(VOORBEELD_CALL1).append("<|im_end|>\n")
                    .append("<|im_start|>user\n<tool_response>\nWrote hello.html. Checks: no problems found.\n</tool_response><|im_end|>\n")
                    .append("<|im_start|>assistant\n").append(VOORBEELD_CALL2).append("<|im_end|>\n")
                    .append(sjabloon.blok("user", "Well done. Now the real task:\n" + taak));
            } else {
                open.append(sjabloon.blok("user", taak));
            }
            open.append("<|im_start|>assistant\n");
            String voorzet = metVoorzet ? VOORZET : "";
            int[] openingIds = tk.encode(open.toString());
            if (openingIds.length + 256 > m.ctx()) {
                System.err.println("de context is te klein voor de opdracht; kies --ctx groter");
                System.exit(1);
            }
            float[] logits = null;
            for (int id : openingIds) logits = m.forward(id);
            laatsteLogits = logits;

            // laatste controleresultaat en inhouds-hash per geschreven bestand:
            // samen de poortwachter van finish én de vastloper-detectie
            Map<String, List<String>> stand = new LinkedHashMap<>();
            Map<String, Integer> vorigeInhoud = new LinkedHashMap<>();
            int herhalingen = 0;
            String voorzetNu = voorzet;

            for (int stap = 1; stap <= stappen; stap++) {
                System.out.printf("%n-- stap %d --%n", stap);
                String antwoord = genereer(m, tk, sampler, voorzetNu);
                if (antwoord == null) { System.out.println("(context vol — gestopt)"); System.exit(1); }

                List<Map<String, Object>> aanroepen = parseToolCalls(antwoord);
                List<String> resultaten = new ArrayList<>();
                boolean zelfdeDezeStap = false;
                if (aanroepen.isEmpty()) {
                    System.out.println("model: " + kort(antwoord.strip()));
                    resultaten.add(laatsteAfgekapt
                            ? "Error: your reply was cut off because it was too long. Write a "
                              + "smaller file and reply with one complete <tool_call> block."
                            : "Error: reply with exactly one <tool_call> block, nothing else.");
                } else {
                    for (Map<String, Object> aanroep : aanroepen) {
                        String resultaat = voerUit(aanroep, map, eisen, stand, vorigeInhoud);
                        if (resultaat == null) {         // finish, en de poort is open
                            System.out.println("\nKLAAR - alle controles groen");
                            System.exit(0);
                        }
                        if (resultaat.startsWith(ZELFDE)) { herhalingen++; zelfdeDezeStap = true; }
                        resultaten.add(resultaat);
                    }
                }
                if (!zelfdeDezeStap) herhalingen = 0;    // alleen een aaneengesloten reeks telt
                for (String r : resultaten) System.out.println("harnas: " + kort(r));

                // vastgelopen in een greedy-herhaling? schud de dobbelsteen, desnoods twee
                // keer steeds harder (herhaalbaar: vaste seed)
                if (herhalingen >= 2 && sampler.temperatuur == 0) {
                    sampler = new Sampler(0.3f, 40, 0.9f, 42);
                    System.out.println("(harnas: tweemaal dezelfde inhoud — temperatuur naar 0,3)");
                } else if (herhalingen >= 4 && sampler.temperatuur < 0.75f) {
                    sampler = new Sampler(0.8f, 40, 0.9f, 42);
                    System.out.println("(harnas: het kopiëren houdt aan — temperatuur naar 0,8)");
                }

                if (voerBinnen(m, tk, sjabloon.toolBeurt(resultaten)) == null) {
                    System.out.println("(context vol — gestopt)");
                    System.exit(1);
                }

                // sturen met de voorzet: staan er problemen open, dan begint het volgende
                // antwoord alvast met een write_file voor het eerste falende bestand
                String slechtPad = stand.entrySet().stream()
                        .filter(e -> !e.getValue().isEmpty()).map(Map.Entry::getKey)
                        .findFirst().orElse(null);
                voorzetNu = (metVoorzet && slechtPad != null)
                        ? "<tool_call>\n{\"name\": \"write_file\", \"arguments\": {\"path\": \""
                          + slechtPad + "\", \"content\": \""
                        : voorzet;
            }
            System.out.println("\nMISLUKT - stappenbudget op zonder groene controles");
            System.exit(1);
        }
    }

    // ---------------------------------------------------------------- de gereedschappen

    static final String ZELFDE = "Error: that is exactly the same content as before.";

    /** Voert één aanroep uit; geeft het resultaat voor het model, of {@code null} bij een geslaagde finish. */
    static String voerUit(Map<String, Object> aanroep, Path map, List<String> eisen,
                          Map<String, List<String>> stand, Map<String, Integer> vorigeInhoud) {
        // kapotte JSON: de positionele parserfout moet het model bereiken, niet een
        // misleidende melding over 'arguments' — quotes vergeten te ontsnappen is bij
        // een klein model de meest voorkomende fout
        if (aanroep.containsKey("fout")) {
            return "Error: invalid JSON inside <tool_call>: " + aanroep.get("fout")
                    + " Escape quotes and backslashes inside strings.";
        }
        Object naam = aanroep.get("name");
        Object argO = aanroep.get("arguments");
        if (!(argO instanceof Map)) return "Error: 'arguments' must be a JSON object.";
        @SuppressWarnings("unchecked") Map<String, Object> arg = (Map<String, Object>) argO;

        if ("write_file".equals(naam)) {
            if (!(arg.get("path") instanceof String pad) || !(arg.get("content") instanceof String inhoud)) {
                return "Error: write_file needs string arguments 'path' and 'content'.";
            }
            System.out.printf("model: write_file %s (%d tekens)%n", pad, inhoud.length());
            String fout = veiligPad(pad, stand.size());
            if (fout != null) return fout;
            if (inhoud.length() > MAX_TEKENS) return "Error: content too large.";
            try {
                Files.writeString(map.resolve(pad), inhoud, StandardCharsets.UTF_8);
            } catch (Exception e) {
                return "Error: could not write file: " + e.getMessage();
            }
            List<String> problemen = isHtml(pad)
                    ? HtmlControles.controleer(inhoud, eisen) : List.of();
            Integer vorige = vorigeInhoud.put(pad, inhoud.hashCode());
            stand.put(pad, problemen);
            if (problemen.isEmpty()) return "Wrote " + pad + ". Checks: no problems found.";
            String kop = vorige != null && vorige == inhoud.hashCode()
                    ? ZELFDE + " It still has " : "Wrote " + pad + ". Checks found ";
            // ontbreekt geëiste letterlijke tekst, geef die dan als Engelse hint mee —
            // getoetst op dezelfde commentaar-gestripte tekst als de controle zelf
            String hint = eisen.stream().filter(e -> !HtmlControles.bevatEis(inhoud, e))
                    .reduce((a, b) -> a + " and " + b)
                    .map(t -> "\nHint: the file must literally contain: " + t)
                    .orElse("");
            return kop + problemen.size() + " problem(s):\n- "
                    + String.join("\n- ", problemen) + hint
                    + "\nFix them by writing the complete file again.";
        }

        if ("finish".equals(naam)) {
            if (stand.isEmpty()) return "Error: nothing was written yet; use write_file first.";
            if (stand.keySet().stream().noneMatch(Agent::isHtml)) {
                return "Error: no HTML file was written yet; the task needs one.";
            }
            List<String> open = new ArrayList<>();
            stand.forEach((pad, problemen) -> problemen.forEach(p -> open.add(pad + ": " + p)));
            if (!open.isEmpty()) {
                return "Error: checks still report problems:\n- " + String.join("\n- ", open)
                        + "\nFix them with write_file before calling finish.";
            }
            System.out.println("model: finish — " + arg.getOrDefault("summary", ""));
            return null;                                 // de poort is open
        }
        return "Error: unknown tool '" + naam + "'. Available: write_file, finish.";
    }

    /** Dezelfde toets overal: de poortwachter mag niet op hoofdletters struikelen. */
    static boolean isHtml(String pad) {
        return pad.toLowerCase(java.util.Locale.ROOT).endsWith(".html");
    }

    /** Binnen de werkmap blijven, en alleen onschuldige bestandssoorten. */
    static String veiligPad(String pad, int aantalBestanden) {
        // het strakke patroon weert ook aanhalingstekens en regeleindes: de sturende
        // voorzet plakt dit pad later ongewijzigd in een JSON-prefix
        if (pad.isBlank() || pad.contains("..") || !pad.matches("[A-Za-z0-9._-]+")) {
            return "Error: path must be a plain file name (letters, digits, dot, dash).";
        }
        String laag = pad.toLowerCase(java.util.Locale.ROOT);
        if (!(laag.endsWith(".html") || laag.endsWith(".css") || laag.endsWith(".js")
                || laag.endsWith(".txt") || laag.endsWith(".md"))) {
            return "Error: only .html, .css, .js, .txt or .md files are allowed.";
        }
        if (aantalBestanden >= MAX_BESTANDEN) return "Error: too many files.";
        return null;
    }

    // ---------------------------------------------------------------- de lus-onderdelen

    private static final Pattern TOOL_CALL =
            Pattern.compile("<tool_call>\\s*(.*?)\\s*</tool_call>", Pattern.DOTALL);

    /** Vist de {@code <tool_call>}-blokken uit een antwoord; ongeldige JSON telt niet mee. */
    static List<Map<String, Object>> parseToolCalls(String antwoord) {
        List<Map<String, Object>> uit = new ArrayList<>();
        Matcher m = TOOL_CALL.matcher(antwoord);
        while (m.find()) {
            try {
                if (Json.parse(m.group(1)) instanceof Map<?, ?> mm) {
                    @SuppressWarnings("unchecked") Map<String, Object> aanroep = (Map<String, Object>) mm;
                    uit.add(aanroep);
                }
            } catch (IllegalArgumentException e) {
                uit.add(Map.of("name", "?", "fout", e.getMessage()));
            }
        }
        return uit;
    }

    /**
     * Genereert één antwoord (greedy) en sluit de beurt af; {@code null} als de context
     * vol is. De voorzet wordt eerst als vaststaande tekst gevoerd: het model "heeft dan
     * al besloten" een gereedschap aan te roepen en hoeft alleen nog in te vullen welk.
     */
    private static String genereer(Motor m, Tokenizer tk, Sampler sampler, String voorzet) {
        StringBuilder sb = new StringBuilder(voorzet);
        float[] logits = voerBinnen(m, tk, voorzet);
        if (logits == null) return null;
        // De protocol-tokens zijn zélf speciale tokens; die horen bij het antwoord en
        // mogen de generatie niet beëindigen zoals <|im_end|> dat doet.
        int toolOpen = tk.idOf("<tool_call>"), toolDicht = tk.idOf("</tool_call>");
        Chat.Utf8Stroom uit = new Chat.Utf8Stroom();
        int volgende = sampler.kies(logits);
        int aantal = 0;
        while (aantal < MAX_ANTWOORD) {
            boolean protocol = volgende == toolOpen || volgende == toolDicht;
            if (!protocol && (volgende == tk.eosId || tk.isSpecial(volgende))) break;
            sb.append(uit.voeg(tk.tokenBytes(volgende)));
            aantal++;
            if (m.position() + 3 > m.ctx()) return null;
            logits = m.forward(volgende);
            volgende = sampler.kies(logits);
        }
        laatsteAfgekapt = aantal >= MAX_ANTWOORD;        // afgekapt is iets anders dan geweigerd
        sb.append(uit.rest());
        for (int id : tk.encode("<|im_end|>\n")) {       // de beurt geldig afsluiten, zoals in Chat
            if (m.position() >= m.ctx()) return null;
            logits = m.forward(id);
        }
        laatsteLogits = logits;
        return sb.toString();
    }

    private static float[] laatsteLogits;
    private static boolean laatsteAfgekapt;

    /** Voert een sjabloonstuk en houdt de logits-stand bij; {@code null} als de context vol is. */
    private static float[] voerBinnen(Motor m, Tokenizer tk, String tekst) {
        int[] ids = tk.encode(tekst);
        if (m.position() + ids.length + 64 > m.ctx()) return null;
        for (int id : ids) laatsteLogits = m.forward(id);
        return laatsteLogits;
    }

    static String kort(String s) {
        String t = s.replace("\n", " ⏎ ");
        return t.length() > 160 ? t.substring(0, 157) + "..." : t;
    }

    private static String eis(String[] args, int i, String vlag) {
        if (i >= args.length) { System.err.println("optie " + vlag + " heeft een waarde nodig"); System.exit(2); }
        return args[i];
    }
}