using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Text;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;

namespace Gott.Services;

/// <summary>
/// Lokaler Embedding-Service (ONNX). Standardmäßig all-MiniLM-L6-v2.
/// Generiert 384-dimensionale Sentence-Embeddings für semantische Suche.
/// Das Modellverzeichnis und die maximale Sequenzlänge sind konfigurierbar.
/// </summary>
public sealed class EmbeddingService : IDisposable
{
    readonly InferenceSession _session;
    readonly Dictionary<string, int> _vocab;
    readonly int _clsId;
    readonly int _sepId;
    readonly int _unkId;
    readonly int _maxSeqLen;
    bool _disposed;

    public bool IsAvailable { get; }
    public int Dimensions => 384;
    public string ModelName { get; }
    public bool CrossEncoderAvailable => false;

    public EmbeddingService(string modelDir = null, int maxSeqLen = 256)
    {
        _maxSeqLen = maxSeqLen;
        if (modelDir == null)
            modelDir = Path.Combine(AppContext.BaseDirectory, "data", "models", "all-MiniLM-L6-v2");
        ModelName = Path.GetFileName(modelDir);
        var modelPath = Path.Combine(modelDir, "model.onnx");
        var vocabPath = Path.Combine(modelDir, "vocab.txt");

        if (!File.Exists(modelPath) || !File.Exists(vocabPath))
        {
            Console.WriteLine("[Embedding] Model not found at " + modelDir + " – service disabled.");
            _vocab = new Dictionary<string, int>();
            IsAvailable = false;
            return;
        }

        try
        {
            var opts = new SessionOptions();
            opts.InterOpNumThreads = 1;
            opts.IntraOpNumThreads = 2;
            opts.GraphOptimizationLevel = GraphOptimizationLevel.ORT_ENABLE_ALL;
            _session = new InferenceSession(modelPath, opts);

            _vocab = new Dictionary<string, int>(50000);
            var lines = File.ReadAllLines(vocabPath);
            for (int i = 0; i < lines.Length; i++)
                _vocab[lines[i]] = i;

            _clsId = _vocab.GetValueOrDefault("[CLS]", 101);
            _sepId = _vocab.GetValueOrDefault("[SEP]", 102);
            _unkId = _vocab.GetValueOrDefault("[UNK]", 100);

            IsAvailable = true;
            Console.WriteLine($"[Embedding] Model loaded. Vocab: {_vocab.Count} tokens.");
        }
        catch (Exception ex)
        {
            Console.WriteLine($"[Embedding] Failed to load model: {ex.Message}");
            _vocab = new Dictionary<string, int>();
            IsAvailable = false;
        }
    }

    /// <summary>Generiert ein 384-dim Embedding für einen Text.</summary>
    public float[] GenerateEmbedding(string text)
    {
        if (!IsAvailable) return Array.Empty<float>();

        var tokenIds = Tokenize(text);
        var seqLen = tokenIds.Count;

        var inputIds = new long[seqLen];
        var attentionMask = new long[seqLen];
        var tokenTypeIds = new long[seqLen];

        for (int i = 0; i < seqLen; i++)
        {
            inputIds[i] = tokenIds[i];
            attentionMask[i] = 1;
            tokenTypeIds[i] = 0;
        }

        var shape = new long[] { 1, seqLen };

        var inputs = new List<NamedOnnxValue>
        {
            NamedOnnxValue.CreateFromTensor("input_ids",
                new DenseTensor<long>(inputIds, new[] { 1, seqLen })),
            NamedOnnxValue.CreateFromTensor("attention_mask",
                new DenseTensor<long>(attentionMask, new[] { 1, seqLen })),
            NamedOnnxValue.CreateFromTensor("token_type_ids",
                new DenseTensor<long>(tokenTypeIds, new[] { 1, seqLen }))
        };

        using var results = _session.Run(inputs);
        var output = results.First().AsTensor<float>();

        // Mean Pooling: Mittelwert über alle Token-Embeddings (attention_mask berücksichtigen)
        var embedding = new float[Dimensions];
        for (int t = 0; t < seqLen; t++)
        {
            for (int d = 0; d < Dimensions; d++)
                embedding[d] += output[0, t, d];
        }
        for (int d = 0; d < Dimensions; d++)
            embedding[d] /= seqLen;

        // L2 Normalize
        Normalize(embedding);

        return embedding;
    }

    /// <summary>WordPiece-Tokenisierung mit BERT-Normalisierung.</summary>
    List<int> Tokenize(string text)
    {
        // BERT Normalization: lowercase, clean control chars
        text = text.ToLowerInvariant();
        var sb = new StringBuilder(text.Length);
        foreach (var c in text)
        {
            if (c == 0 || c == 0xFFFD || char.IsControl(c)) continue;
            if (IsChineseChar(c)) { sb.Append(' '); sb.Append(c); sb.Append(' '); }
            else sb.Append(c);
        }
        text = sb.ToString();

        // BERT Pre-tokenize: split on whitespace and punctuation
        var words = BertPreTokenize(text);

        var tokenIds = new List<int> { _clsId };
        int maxTokens = _maxSeqLen - 2; // Reserve [CLS] and [SEP]

        foreach (var word in words)
        {
            if (tokenIds.Count - 1 >= maxTokens) break;

            var subTokens = WordPieceTokenize(word);
            foreach (var st in subTokens)
            {
                if (tokenIds.Count - 1 >= maxTokens) break;
                tokenIds.Add(st);
            }
        }

        tokenIds.Add(_sepId);
        return tokenIds;
    }

    List<int> WordPieceTokenize(string word)
    {
        if (word.Length > 100)
            return new List<int> { _unkId };

        var tokens = new List<int>();
        int start = 0;

        while (start < word.Length)
        {
            int end = word.Length;
            int foundId = -1;

            while (start < end)
            {
                var substr = word.Substring(start, end - start);
                if (start > 0) substr = "##" + substr;

                if (_vocab.TryGetValue(substr, out var id))
                {
                    foundId = id;
                    break;
                }
                end--;
            }

            if (foundId < 0)
            {
                tokens.Add(_unkId);
                break;
            }

            tokens.Add(foundId);
            start = end;
        }

        return tokens;
    }

    static List<string> BertPreTokenize(string text)
    {
        var words = new List<string>();
        var current = new StringBuilder();

        foreach (var c in text)
        {
            if (char.IsWhiteSpace(c))
            {
                if (current.Length > 0) { words.Add(current.ToString()); current.Clear(); }
            }
            else if (IsPunctuation(c))
            {
                if (current.Length > 0) { words.Add(current.ToString()); current.Clear(); }
                words.Add(c.ToString());
            }
            else
            {
                current.Append(c);
            }
        }

        if (current.Length > 0) words.Add(current.ToString());
        return words;
    }

    static bool IsPunctuation(char c)
    {
        var cat = char.GetUnicodeCategory(c);
        return cat == UnicodeCategory.ConnectorPunctuation ||
               cat == UnicodeCategory.DashPunctuation ||
               cat == UnicodeCategory.OpenPunctuation ||
               cat == UnicodeCategory.ClosePunctuation ||
               cat == UnicodeCategory.InitialQuotePunctuation ||
               cat == UnicodeCategory.FinalQuotePunctuation ||
               cat == UnicodeCategory.OtherPunctuation;
    }

    static bool IsChineseChar(char c)
    {
        int cp = c;
        return (cp >= 0x4E00 && cp <= 0x9FFF) || (cp >= 0x3400 && cp <= 0x4DBF) ||
               (cp >= 0x20000 && cp <= 0x2A6DF) || (cp >= 0x2A700 && cp <= 0x2B73F) ||
               (cp >= 0x2B740 && cp <= 0x2B81F) || (cp >= 0x2B820 && cp <= 0x2CEAF) ||
               (cp >= 0xF900 && cp <= 0xFAFF) || (cp >= 0x2F800 && cp <= 0x2FA1F);
    }

    static void Normalize(float[] vec)
    {
        double sum = 0;
        for (int i = 0; i < vec.Length; i++) sum += vec[i] * vec[i];
        var norm = (float)Math.Sqrt(sum);
        if (norm > 0)
            for (int i = 0; i < vec.Length; i++) vec[i] /= norm;
    }

    /// <summary>Cosine Similarity zwischen zwei normalisierten Vektoren (= Dot Product).</summary>
    public static float CosineSimilarity(float[] a, float[] b)
    {
        if (a.Length != b.Length) return 0;
        float dot = 0;
        for (int i = 0; i < a.Length; i++) dot += a[i] * b[i];
        return dot;
    }

    /// <summary>Serialisiert float[] zu byte[] für DB-Speicherung.</summary>
    public static byte[] ToBytes(float[] vec)
    {
        var bytes = new byte[vec.Length * 4];
        Buffer.BlockCopy(vec, 0, bytes, 0, bytes.Length);
        return bytes;
    }

    /// <summary>Deserialisiert byte[] aus DB zu float[].</summary>
    public static float[] FromBytes(byte[] bytes)
    {
        var vec = new float[bytes.Length / 4];
        Buffer.BlockCopy(bytes, 0, vec, 0, bytes.Length);
        return vec;
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            _session?.Dispose();
            _disposed = true;
        }
    }
}