using System;
using System.Threading.Tasks;
using LightRAG.Core.Abstractions;
using Unity.InferenceEngine;
using SentisModels;
using OnDeviceAgent.Inference;

namespace OnDeviceAgent.AgentCore
{
// Framework adapter over SentisModels.E5Embedder: adds Android-safe loading, ITextEmbedder, and the LightRAG EmbeddingFunc bridge.
public sealed class E5TextEmbedder : ITextEmbedder
{
    public const int FeatureDim = E5Embedder.FeatureDim;
    const int MaxTokens = 512;            // e5 context limit (used for the EmbeddingFunc token budget)

    readonly E5Embedder m_Core;

    public bool Ready => m_Core.Ready;

    public E5TextEmbedder(BackendType backend = BackendType.CPU, Action<string> log = null)
    {
        m_Core = new E5Embedder(backend, log);
    }

    public bool Load(string modelPath, string tokenizerPath)
    {
        var root = (System.IO.Path.GetDirectoryName(modelPath) ?? string.Empty).Replace('\\', '/');
        return m_Core.Load(
            ModelRootProvisioner.EnsureModelRoot(root),
            System.IO.Path.GetFileName(modelPath),
            System.IO.Path.GetFileName(tokenizerPath));
    }

    public float[] EmbedText(string text, bool isQuery = false) => m_Core.EmbedText(text, isQuery);

    public Task<float[]> EmbedTextAsync(string text, bool isQuery = false) => m_Core.EmbedTextAsync(text, isQuery);

    public EmbeddingFunc AsEmbeddingFunc(IUnityMainThreadDispatcher dispatcher)
    {
        if (dispatcher == null) throw new ArgumentNullException(nameof(dispatcher));
        // Sync EmbedText, not async: async readback can hang the serialized queue in batchmode (no SynchronizationContext).
        return new EmbeddingFunc(
            FeatureDim,
            (texts, context, ct) => dispatcher.RunOnMainAsync<float[][]>(() =>
            {
                var isQuery = context == "query";
                var vectors = new float[texts.Count][];
                for (var i = 0; i < texts.Count; i++)
                {
                    ct.ThrowIfCancellationRequested();
                    vectors[i] = EmbedText(texts[i], isQuery);
                }
                return Task.FromResult(vectors);
            }),
            MaxTokenSize: MaxTokens,
            ModelName: "multilingual-e5-small",
            SupportsAsymmetric: true);
    }

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