package com.sd.hmct_small_model

import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer
import android.content.Context
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import org.json.JSONException
import org.json.JSONObject
import org.pytorch.IValue
import org.pytorch.Module
import org.pytorch.Tensor
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.io.OutputStream
import java.nio.file.Paths

class InferenceAppK(var basePath: String = "models") {
    private var intentModel: Module? = null
    private var slotModel: Module? = null
    private lateinit var intentTokenizer: HuggingFaceTokenizer
    private lateinit var slotTokenizer: HuggingFaceTokenizer
    private var intentLabels = mutableListOf<String>()
    private var slotLabels = mutableListOf<String>()

    init {

    }


    suspend fun initialize(context: Context) {
        try {
            var parentFileDir = File(context.filesDir, "$basePath/intent_tokenizer")
            if (!parentFileDir.exists())
                parentFileDir.mkdirs()
            parentFileDir = File(context.filesDir, "$basePath/slot_tokenizer")
            if (!parentFileDir.exists())
                parentFileDir.mkdir()


            var path = assetFilePath(context, basePath.plus("/intent_model_traced.pt"))
            intentModel = Module.load(path)
             path = assetFilePath(context, basePath.plus("/slot_model_traced.pt"))
            slotModel = Module.load(path)

            path = assetFilePath(context, "$basePath/intent_tokenizer/tokenizer.json")
            intentTokenizer = HuggingFaceTokenizer.newInstance(Paths.get(path))
            path = assetFilePath(context, "$basePath/slot_tokenizer/tokenizer.json")
            slotTokenizer = HuggingFaceTokenizer.newInstance(Paths.get(path))

//            intentLabels =
//                Utils.fromJson(Paths.get(basePath + "/intent_labels.json"), MutableList::class.java)
//            slotLabels =
//                Utils.fromJson(Paths.get(basePath + "/slot_labels.json"), MutableList::class.java)
        } catch (e: IOException) {
            e.printStackTrace()
        }
    }

    suspend fun assetFilePath(context: Context, assetName: String): String {
        return withContext(Dispatchers.IO) {
            val file = File(context.filesDir, assetName)
            if (file.exists() && file.length() > 0) {
                return@withContext file.absolutePath
            }
            try {
                var inputs = context.assets.open(assetName)
                var os: OutputStream = FileOutputStream(file)
                val buffer = ByteArray(4 * 1024)
                var read: Int
                while ((inputs.read(buffer).also { read = it }) != -1) {
                    os.write(buffer, 0, read)
                }
                os.flush()
                return@withContext file.absolutePath
            } catch (e: IOException) {
                e.printStackTrace()
                return@withContext ""
            }
        }
    }


    fun predict(text: String): JSONObject {
        val encIntent = intentTokenizer.encode(text)
        val ids = encIntent.getIds()
        val mask = encIntent.getAttentionMask()
        val inputIds = Tensor.fromBlob(ids, longArrayOf(1, ids.size.toLong()))
        val attentionMask = Tensor.fromBlob(mask, longArrayOf(1, mask.size.toLong()))
        val logitsIntent =
            intentModel!!.forward(IValue.from(inputIds), IValue.from(attentionMask)).toTensor()

        val intentOut = logitsIntent.getDataAsFloatArray()
        val intentIdx = argMax(intentOut)
        val intent = intentLabels.get(intentIdx)

        val encSlot = slotTokenizer.encode(text)
        val sIds = Tensor.fromBlob(encSlot.getIds(), longArrayOf(1, encSlot.getIds().size.toLong()))
        val sMask = Tensor.fromBlob(
            encSlot.getAttentionMask(),
            longArrayOf(1, encSlot.getAttentionMask().size.toLong())
        )
        val logitsSlot = slotModel?.forward(IValue.from(sIds), IValue.from(sMask))?.toTensor()
//        val slotOut: LongArray = logitsSlot.argmax(2).getDataAsLongArray()
        val slotOut = logitsSlot!!.getDataAsLongArray()

//        val params = extractParams(encSlot.getTokens(), encSlot.getOffsets(), slotOut)
        val params = extractParams(encSlot.getTokens(), arrayOf(), slotOut)
        val resp = JSONObject()
        resp.put("text", text)
        resp.put("intent", intent)
        resp.put("params", params)
        return resp
    }

    @Throws(JSONException::class)
    private fun extractParams(
        tokens: Array<String?>,
        offsets: Array<LongArray?>,
        preds: LongArray
    ): JSONObject {
        val obj = JSONObject()
        val n = preds.size
        var curSlot: String? = null
        var curText: String? = ""
        for (i in 0..<n) {
            val lbl = slotLabels.get(preds[i].toInt())
            if (lbl == "O" || !lbl.startsWith("B-") && !lbl.startsWith("I-")) {
                if (curSlot != null) {
                    obj.put(curSlot, curText)
                    curSlot = null
                    curText = ""
                }
            } else {
                val parts = lbl.split("-".toRegex(), limit = 2).toTypedArray()
                val prefix = parts[0]
                val slot = parts[1]
                if (prefix == "B") {
                    if (curSlot != null) {
                        obj.put(curSlot, curText)
                    }
                    curSlot = slot
                    curText = tokenToText(tokens[i]!!, offsets[i])
                } else if (prefix == "I" && slot == curSlot) {
                    curText += tokenToText(tokens[i]!!, offsets[i])
                }
            }
        }
        if (curSlot != null) obj.put(curSlot, curText)
        return obj
    }

    private fun tokenToText(token: String, offs: LongArray?): String {
        var t = token
        if (t.startsWith("##") || t.startsWith("▁")) t = t.substring(1)
        return t
    }

    private fun argMax(arr: FloatArray): Int {
        var idx = 0
        var mx = arr[0]
        for (i in 1..<arr.size) if (arr[i] > mx) {
            mx = arr[i]
            idx = i
        }
        return idx
    }
}
