using CloudBuilder.Core.DatabaseAccessor.Entity;
using CloudBuilder.Core.Service;
using DocumentFormat.OpenXml.Wordprocessing;
using EnvDTE80;
using Irony.Parsing;
using Microsoft.ML;
using Microsoft.ML.Data;
using Microsoft.ML.Transforms.Onnx;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace CloudBuilder.AI.Service.Utility
{
    public class ChineseNameExtractor
    {
        private BertChineseTokenizer tokenizer;
        private OnnxScoringEstimator pipeline;
        private PredictionEngine<NerInput, NerOutput> engine;
        public ChineseNameExtractor(IApplicationService service)
        {
            InitMLContext(service);
        }

        private void InitMLContext(IApplicationService service)
        {
            // 1. 初始化 MLContext
            var mlContext = new MLContext();

            string modelDir = service.Configuration["FileServiceSettings:OnnxDirectory"];
            // 2. 检查必要文件
            string modelPath = Path.Combine(modelDir, "model.onnx");
            string vocabPath = Path.Combine(modelDir, "vocab.txt");

            if (!File.Exists(modelPath) || !File.Exists(vocabPath))
            {
                throw new ValidatedException("请确保 model.onnx 和 vocab.txt 在程序目录下。");
            }

            // 3. 加载词汇表与标签映射
            tokenizer = new BertChineseTokenizer(vocabPath);

            // 4. 定义模型输入输出的精确形状
            //    通过 Netron (https://netron.app/) 查看模型结构后，你可以确认这些形状。
            var shapeDictionary = new Dictionary<string, int[]>
                {
                    { "input_ids", new[] { 1, 128 } },
                    { "attention_mask", new[] { 1, 128 } },
                    { "token_type_ids", new[] { 1, 128 } },
                    { "logits", new[] { 1, 128, 29 } }   // 29 个标签
            };

            // 5. 构建管道（使用空数据视图和 shapeDictionary）
            var emptyData = mlContext.Data.LoadFromEnumerable(new List<NerInput>());
            pipeline = mlContext.Transforms.ApplyOnnxModel(
               modelFile: modelPath,
               inputColumnNames: new[] { "input_ids", "attention_mask", "token_type_ids" },
               outputColumnNames: new[] { "logits" },
               shapeDictionary: shapeDictionary,
               gpuDeviceId: null,        // 若要用 GPU，可改为 0 等
               fallbackToCpu: true
           );

            // 6. 拟合并创建预测引擎
            //Console.WriteLine("正在加载 ONNX 模型...");
            var transformer = pipeline.Fit(emptyData);
            engine = mlContext.Model.CreatePredictionEngine<NerInput, NerOutput>(transformer);
            //Console.WriteLine("模型加载完成。");
        }

        public List<string> ExtractNames(string text)
        {
            var tokenized = tokenizer.Tokenize(text);

            var input = new NerInput
            {
                InputIds = tokenized.InputIds,
                AttentionMask = tokenized.AttentionMask,
                TokenTypeIds = tokenized.TokenTypeIds
            };

            var prediction = engine.Predict(input);

            return ExtractNames(prediction, tokenized, tokenizer.IdToLabel);
        }
        /// <summary>
        /// 根据模型输出的 logits 和分词信息，提取出人名（BIO 格式）。
        /// </summary>
        /// <param name="output">模型预测输出。</param>
        /// <param name="tokenized">分词后的结果。</param>
        /// <param name="idToLabel">标签ID到名称的映射数组。</param>
        /// <returns>识别出的人名列表。</returns>
        private List<string> ExtractNames(NerOutput output, TokenizedResult tokenized, string[] idToLabel)
        {
            var names = new List<string>();
            string currentName = "";
            int seqLen = 128;
            int numLabels = idToLabel.Length;

            for (int i = 0; i < seqLen; i++)
            {
                // 跳过填充 token（attention_mask 为 0）
                if (tokenized.AttentionMask[i] == 0)
                    continue;

                int startIdx = i * numLabels;
                float maxVal = float.MinValue;
                int maxIdx = 0;
                for (int j = 0; j < numLabels; j++)
                {
                    float val = output.Logits[startIdx + j];
                    if (val > maxVal)
                    {
                        maxVal = val;
                        maxIdx = j;
                    }
                }

                string label = idToLabel[maxIdx];
                string token = tokenized.Tokens[i];

                if (label == "B-PER")
                {
                    if (currentName != "")
                        names.Add(currentName);
                    currentName = token;
                }
                else if (label == "I-PER")
                {
                    currentName += token;
                }
                else
                {
                    if (currentName != "")
                    {
                        names.Add(currentName);
                        currentName = "";
                    }
                }
            }

            if (currentName != "")
                names.Add(currentName);
            return names;
        }

    }


    // ========== 数据模型定义 ==========
    // 以下类的 VectorType 和 ColumnName 必须与 shapeDictionary 及 ONNX 模型定义完全一致。

    /// <summary>
    /// ONNX 模型的输入。假设 batch=1，序列长度=128，三个 int64 张量。
    /// </summary>
    public class NerInput
    {
        [VectorType(1, 128)]
        [ColumnName("input_ids")]
        public long[] InputIds { get; set; }

        [VectorType(1, 128)]
        [ColumnName("attention_mask")]
        public long[] AttentionMask { get; set; }

        [VectorType(1, 128)]
        [ColumnName("token_type_ids")]
        public long[] TokenTypeIds { get; set; }
    }

    /// <summary>
    /// ONNX 模型的输出。假设 logits 形状为 [1, 128, num_labels]。
    /// </summary>
    public class NerOutput
    {
        [VectorType(1, 128, 29)]   // 与 shapeDictionary 一致
        [ColumnName("logits")]
        public float[] Logits { get; set; }
    }

    // ========== 分词器与词汇表加载 ==========

    /// <summary>
    /// 简易的 BERT 中文分词器，按字进行分割，适用于中文 NER 任务。
    /// 特殊 token：[CLS] = 101, [SEP] = 102, [PAD] = 0。
    /// </summary>
    public class BertChineseTokenizer
    {
        private readonly Dictionary<string, int> _tokenToId;
        /// <summary>
        /// 标签ID到名称的映射，其顺序和内容必须与模型输出完全一致。
        /// 此处的映射基于 bert-base-chinese-finetuned-ner 模型常见的标签体系。
        /// </summary>
        public readonly string[] IdToLabel;

        public BertChineseTokenizer(string vocabPath, int maxSeqLength = 128)
        {
            // 加载词汇表
            _tokenToId = new Dictionary<string, int>();
            var lines = File.ReadAllLines(vocabPath);
            for (int i = 0; i < lines.Length; i++)
                _tokenToId[lines[i]] = i;

            // 默认标签映射，顺序务必与模型输出一致
            IdToLabel = new string[] {
    "O",           // 0
    "B-ORG",       // 1
    "I-ORG",       // 2
    "B-PER",       // 3
    "I-PER",       // 4
    "B-TIME",      // 5
    "I-TIME",      // 6
    "B-LOC",       // 7
    "I-LOC",       // 8
    "B-POSITION",  // 9
    "I-POSITION",  // 10
    "B-COMPANY",   // 11
    "I-COMPANY",   // 12
    "B-GAME",      // 13
    "I-GAME",      // 14
    "B-GOVERNMENT",// 15
    "I-GOVERNMENT",// 16
    "B-SCENE",     // 17
    "I-SCENE",     // 18
    "B-SUBJECT",   // 19
    "I-SUBJECT",   // 20
    "B-CREATION",  // 21
    "I-CREATION",  // 22
    "B-FOOD",      // 23
    "I-FOOD",      // 24
    "B-MOVIE",     // 25
    "I-MOVIE",     // 26
    "[CLS]",       // 27
    "[SEP]"        // 28
};
        }

        /// <summary>
        /// 对中文句子进行分词，并生成模型所需的 input_ids 等。
        /// </summary>
        /// <param name="text">输入的中文句子。</param>
        /// <returns>分词后的结果，包括 token IDs 和 tokens 本身。</returns>
        public TokenizedResult Tokenize(string text)
        {
            const int maxLen = 128;
            var inputIds = new long[maxLen];
            var attentionMask = new long[maxLen];
            var tokenTypeIds = new long[maxLen];
            var tokens = new string[maxLen];

            // 将所有字符当作 token（BERT 中文模型常用）
            var chars = text.ToCharArray();
            int charIndex = 0;

            for (int i = 0; i < maxLen; i++)
            {
                if (i == 0)  // [CLS]
                {
                    inputIds[i] = 101;
                    tokens[i] = "[CLS]";
                }
                else if (i == chars.Length + 1)  // [SEP]
                {
                    inputIds[i] = 102;
                    tokens[i] = "[SEP]";
                }
                else if (i > chars.Length + 1)  // padding
                {
                    inputIds[i] = 0;
                    tokens[i] = "[PAD]";
                }
                else  // 实际字符
                {
                    string c = chars[charIndex].ToString();
                    charIndex++;
                    if (_tokenToId.TryGetValue(c, out int id))
                    {
                        inputIds[i] = id;
                        tokens[i] = c;
                    }
                    else
                    {
                        // 如果词汇表没有这个字，用 [UNK] token
                        inputIds[i] = _tokenToId["[UNK]"];
                        tokens[i] = "[UNK]";
                    }
                }

                // attention_mask：有效 token 为 1，padding 为 0
                attentionMask[i] = (i <= chars.Length + 1) ? 1 : 0;
                tokenTypeIds[i] = 0; // 单句都为 0
            }

            return new TokenizedResult
            {
                InputIds = inputIds,
                AttentionMask = attentionMask,
                TokenTypeIds = tokenTypeIds,
                Tokens = tokens
            };
        }
    }

    /// <summary>
    /// 存储分词后的结果。
    /// </summary>
    public class TokenizedResult
    {
        public long[] InputIds { get; set; }
        public long[] AttentionMask { get; set; }
        public long[] TokenTypeIds { get; set; }
        public string[] Tokens { get; set; }
    }
}