// elephant-1 ONNX demo: chat generation + one end-to-end scored decision.
//
//   dotnet run -- --model ../model-cuda-int4 --cuda --prompt "What is the capital of France?"
//   dotnet run -- --model ../model-cuda-int4 --cuda --score-demo
//
// Tokenizer: gigatoken_cs on the build's tokenizer.json — verified byte-identical
// to HuggingFace transformers on 19 texts (unicode, CJK, emoji, specials).
// Chat bytes below mirror the checkpoint's chat template for a single user turn
// (verified against apply_chat_template output).
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;

static class Program
{
    const string SystemPrompt =
        "Apply the question to the state. Choose exactly one of the listed options. " +
        "Respond with only its uppercase letter, with no explanation or reasoning.";

    static int Main(string[] args)
    {
        string modelDir = Arg(args, "--model", "../model-cuda-int4");
        bool cuda = Has(args, "--cuda");
        string prompt = Arg(args, "--prompt", "What is the capital of France? Answer with the city name only.");
        bool scoreDemo = Has(args, "--score-demo");
        int maxNewTokens = int.Parse(Arg(args, "--max-new-tokens", "512"));
        // The build dirs carry the GenAI-fixed tokenizer (ByteLevel-only), which
        // gigatoken_cs 0.12.0 cannot load; the pristine original tokenizes
        // byte-identically (verified 19/19), so load that instead.
        string tokPath = Arg(args, "--tokenizer", "../orig_tokenizer/tokenizer.json");

        var tok = Gigatoken.Tokenizer.FromJson(File.ReadAllText(tokPath));
        using (tok)
        using (var lm = new OnnxLm(Path.Combine(modelDir, "model.onnx"), cuda))
        {
            if (scoreDemo) return ScoreDemo(tok, lm);
            Generate(tok, lm, modelDir, prompt, maxNewTokens);
            return 0;
        }
    }

    // ---- chat generation (causal decode on the bidi graph) ----
    static void Generate(Gigatoken.Tokenizer tok, OnnxLm lm, string modelDir, string prompt, int maxNewTokens)
    {
        string chat = "<|im_start|>user\n" + prompt + "<|im_end|>\n" +
                      "<|im_start|>assistant\n<think>\n\n</think>\n\n";
        long[] ids = tok.Encode(chat).Select(i => (long)i).ToArray();
        Console.WriteLine($"  prompt tokens: {ids.Length}");
        var eos = EosIds(modelDir);
        var t0 = DateTime.UtcNow;
        var outIds = lm.Generate(ids, maxNewTokens, eos, t =>
        {
            Console.Write(tok.DecodeToString(new[] { (uint)t }));
            Console.Out.Flush();
        });
        Console.WriteLine($"\n  tokens={outIds.Count} total={(DateTime.UtcNow - t0).TotalSeconds:F1}s");
    }

    static HashSet<long> EosIds(string modelDir)
    {
        var set = new HashSet<long>();
        try
        {
            var doc = JsonDocument.Parse(File.ReadAllText(Path.Combine(modelDir, "genai_config.json")));
            var eos = doc.RootElement.GetProperty("model").GetProperty("eos_token_id");
            if (eos.ValueKind == JsonValueKind.Array)
                foreach (var e in eos.EnumerateArray()) set.Add(e.GetInt64());
            else set.Add(eos.GetInt64());
        }
        catch { set.Add(248046); set.Add(248044); }
        return set;
    }

    // ---- one scored decision: support.json's refund question ----
    // Prompt rendering mirrors score_elephant.py (qev-labels-v2) for a
    // criterion-less noul question; the program prints the prompt SHA256 so
    // you can check byte-equality against the Python renderer.
    static int ScoreDemo(Gigatoken.Tokenizer tok, OnnxLm lm)
    {
        string state = "Customer: my order arrived broken, I want my money back.";
        string instructions = "Should the agent offer a refund?";
        string t = "the statement is true / the answer is yes";
        string f = "the statement is false / the answer is no";
        string suffix = $"Question: {instructions}\n\nOptions:\nA. yes: {t}\nB. no: {f}\n\nAnswer with one letter: A, B.";
        string user = $"<state>\n{state}\n</state>\n\n" + suffix;
        string full = "<|im_start|>system\n" + SystemPrompt + "<|im_end|>\n" +
                      "<|im_start|>user\n" + user + "<|im_end|>\n" +
                      "<|im_start|>assistant\n<think>\n\n</think>\n\n";
        Console.WriteLine("prompt sha256: " + Sha(full));

        long[] ids = tok.Encode(full).Select(i => (long)i).ToArray();
        int prefix = CommonPrefix(tok, full, $"<state>\n{state}\n</state>\n\n");
        Console.WriteLine($"branch tokens: {ids.Length}, prefix tokens: {prefix}");

        var bias = new float[ids.Length * ids.Length];
        for (int i = 0; i < ids.Length; i++)
            for (int j = 0; j < ids.Length; j++)
                bias[i * ids.Length + j] = (j < prefix || i >= prefix) ? 0f : OnnxLm.NegInf;

        uint[] slots = { tok.Encode("A")[0], tok.Encode("B")[0] };
        if (tok.Encode("A").Length != 1 || tok.Encode("B").Length != 1)
        { Console.Error.WriteLine("slot labels are not single tokens"); return 1; }

        var (last, st) = lm.Prefill(ids, bias);
        using (st)
        {
            double[] z = { last[slots[0]], last[slots[1]] };
            double m = Math.Max(z[0], z[1]);
            double[] p = { Math.Exp(z[0] - m), Math.Exp(z[1] - m) };
            double s = p[0] + p[1];
            Console.WriteLine($"refund noul: {p[0] / s:F4} (Python: 0.8791 CUDA int8, 0.9342 CUDA int4, 0.9069 CPU, 0.8921 CPU int8)");
            return 0;
        }
    }

    // Shared-prefix length: tokenize every branch fully, then find the longest
    // common token prefix with the state-block text (marker method, mirrors
    // qev _prefix_and_suffix_ids: the shared prefix ends at the state block).
    static int CommonPrefix(Gigatoken.Tokenizer tok, string full, string stateBlock)
    {
        long[] fullIds = tok.Encode(full).Select(i => (long)i).ToArray();
        const string mark = "QEV_PREFIX_MARK";
        string head = "<|im_start|>system\n" + SystemPrompt + "<|im_end|>\n" +
                      "<|im_start|>user\n" + stateBlock + mark;
        long[] headIds = tok.Encode(head).Select(i => (long)i).ToArray();
        int n = Math.Min(headIds.Length, fullIds.Length);
        while (n > 0 && !fullIds.Take(n).SequenceEqual(headIds.Take(n))) n--;
        return n;
    }

    static string Sha(string s)
    {
        var h = SHA256.HashData(Encoding.UTF8.GetBytes(s));
        return BitConverter.ToString(h).Replace("-", "").ToLowerInvariant();
    }

    static string Arg(string[] a, string name, string dflt)
    {
        for (int i = 0; i + 1 < a.Length; i++) if (a[i] == name) return a[i + 1];
        return dflt;
    }

    static bool Has(string[] a, string name) => a.Contains(name);
}