using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;

namespace Backend;

/// <summary>
/// Loads a PPO policy ONNX model once and provides deterministic action selection.
/// Thread-safe: <see cref="InferenceSession"/> is safe for concurrent Run() calls.
/// </summary>
public sealed class PolicyRunner : IDisposable
{
    public const int ObsDim = 7;
    public const int NActions = 4;

    private readonly InferenceSession _session;
    private readonly string _inputName;

    public PolicyRunner(string onnxPath)
    {
        if (!File.Exists(onnxPath))
            throw new FileNotFoundException($"ONNX model not found at {onnxPath}");
        _session = new InferenceSession(onnxPath);
        _inputName = _session.InputMetadata.Keys.First();
    }

    /// <summary>Run one inference and return the argmax action index.</summary>
    public int SelectAction(ReadOnlySpan<float> obs)
    {
        if (obs.Length != ObsDim)
            throw new ArgumentException($"expected obs of length {ObsDim}, got {obs.Length}", nameof(obs));

        var tensor = new DenseTensor<float>(new[] { 1, ObsDim });
        for (int i = 0; i < ObsDim; i++) tensor[0, i] = obs[i];

        using var results = _session.Run(new[]
        {
            NamedOnnxValue.CreateFromTensor(_inputName, tensor)
        });

        var logits = results.First().AsEnumerable<float>().ToArray();
        // argmax
        int best = 0;
        float bestVal = logits[0];
        for (int i = 1; i < logits.Length; i++)
        {
            if (logits[i] > bestVal) { best = i; bestVal = logits[i]; }
        }
        return best;
    }

    public void Dispose() => _session.Dispose();
}