using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using Models;
using ToolsServices;

namespace Reranker;

public static class BgeReranker
{
    private static InferenceSession? _session;
    private static Tokenizer? _tokenizer;

    private static bool LoadModeles()
    {
        LoggerService.LogDebug("BgeReranker.LoadModeles");
        try
        {
            LoggerService.LogDebug(typeof(InferenceSession).Assembly.Location);
            var pathModele = Path.GetDirectoryName(System.Reflection.Assembly.GetExecutingAssembly().Location);
            var modeleFullPath = Path.Combine(pathModele!, "onnx");
            _tokenizer = new Tokenizer(Path.Combine(modeleFullPath, "tokenizer.json"));
            /*
            var opts = new SessionOptions();
            opts.AppendExecutionProvider_CPU(); // forcé
            var session = new InferenceSession("onnx/model.onnx", opts);
            */
            _session = new InferenceSession(Path.Combine(modeleFullPath, "model.onnx"));
            return true;
        }
        catch (Exception ex)
        {
            LoggerService.LogError($"BgeReranker.LoadModeles : {ex.Message}");
            throw;

        }
    }

    public static List<RankedDocument> Rerank(string query, List<SearchResult> docs, int topK = 5)
    {
        LoggerService.LogDebug("BgeReranker.Rerank");
        if (_tokenizer == null || _session == null)
        {
            var b = LoadModeles();
            if(!b)
            {
                return new List<RankedDocument>();
            }
        }
        var scored = new List<RankedDocument>();

        foreach (var doc in docs)
        {
            float score = ScorePair(query, doc.Text);
            scored.Add(new RankedDocument
            {
                Text = doc.Text,
                Nom_Fichier = doc.Nom_Fichier,
                Score = score
            });
        }

        return scored
            .OrderByDescending(x => x.Score)
            .Take(topK)
            .ToList();
    }

    private static float ScorePair(string query, string document)
    {
        LoggerService.LogDebug("BgeReranker.ScorePair");
        var encoded = _tokenizer!.EncodePair(query, document, 100);

        var inputIds = new DenseTensor<long>(new[] { 1, encoded.inputIds.Length });
        var attentionMask = new DenseTensor<long>(new[] { 1, encoded.attentionMask.Length });

        for (int i = 0; i < encoded.inputIds.Length; i++)
        {
            inputIds[0, i] = encoded.inputIds[i];
            attentionMask[0, i] = encoded.attentionMask[i];
        }

        var inputs = new List<NamedOnnxValue>
        {
            NamedOnnxValue.CreateFromTensor("input_ids", inputIds),
            NamedOnnxValue.CreateFromTensor("attention_mask", attentionMask)
        };

        using var result = _session!.Run(inputs);
        var output = result.First().AsEnumerable<float>().First();

        return output;
    }
}
