import os
import gradio as gr
import numpy as np
import onnxruntime as ort
import sentencepiece as spm
from huggingface_hub import snapshot_download


# =========================
# DOWNLOAD MODEL
# =========================

MODEL_ID = "Phitran21/AnViMt"

print("Đang tải AnViMt...")

MODEL_DIR = snapshot_download(
    repo_id=MODEL_ID,
    repo_type="model",
)

print("Model directory:", MODEL_DIR)


# =========================
# PATH
# =========================

ENCODER_PATH = os.path.join(
    MODEL_DIR,
    "encoder_model.onnx"
)

DECODER_PATH = os.path.join(
    MODEL_DIR,
    "decoder_model.onnx"
)

SPM_PATH = os.path.join(
    MODEL_DIR,
    "spm.model"
)


# =========================
# SETTINGS
# =========================

MAX_INPUT_LEN = 256
MAX_OUTPUT_LEN = 256

PAD_ID = 0
UNK_ID = 1
BOS_ID = 2
EOS_ID = 3


# =========================
# CHECK FILES
# =========================

for path in [
    ENCODER_PATH,
    DECODER_PATH,
    SPM_PATH,
]:
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"Không tìm thấy file: {path}"
        )


# =========================
# LOAD TOKENIZER
# =========================

print("Đang load tokenizer...")

sp = spm.SentencePieceProcessor()
sp.load(SPM_PATH)

print(
    f"Tokenizer loaded | "
    f"vocab={sp.get_piece_size()} | "
    f"PAD={sp.pad_id()} "
    f"BOS={sp.bos_id()} "
    f"EOS={sp.eos_id()}"
)


# =========================
# LOAD ONNX MODEL
# =========================

print("Đang load ONNX model...")

providers = [
    "CPUExecutionProvider"
]

encoder_session = ort.InferenceSession(
    ENCODER_PATH,
    providers=providers
)

decoder_session = ort.InferenceSession(
    DECODER_PATH,
    providers=providers
)

print("AnViMt ready!")


# =========================
# TRANSLATE
# =========================

def translate(text):

    text = text.strip()

    if not text:
        return ""

    # -------------------------
    # ENCODE INPUT
    # -------------------------

    src_ids = sp.encode(
        text,
        out_type=int
    )

    src_ids = src_ids[
        :MAX_INPUT_LEN - 2
    ]

    src_ids = (
        [BOS_ID]
        + src_ids
        + [EOS_ID]
    )

    input_ids = np.array(
        [src_ids],
        dtype=np.int64
    )

    attention_mask = np.ones_like(
        input_ids,
        dtype=np.int64
    )


    # -------------------------
    # ENCODER
    # -------------------------

    encoder_outputs = encoder_session.run(
        None,
        {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }
    )

    encoder_hidden_states = encoder_outputs[0]


    # -------------------------
    # DECODER GREEDY
    # -------------------------

    generated = [BOS_ID]

    for _ in range(MAX_OUTPUT_LEN):

        decoder_input_ids = np.array(
            [generated],
            dtype=np.int64
        )

        decoder_inputs = {}

        for inp in decoder_session.get_inputs():

            name = inp.name

            if name == "input_ids":

                decoder_inputs[name] = (
                    decoder_input_ids
                )

            elif (
                "encoder_hidden_states"
                in name
            ):

                decoder_inputs[name] = (
                    encoder_hidden_states
                )

            elif (
                "encoder_attention_mask"
                in name
            ):

                decoder_inputs[name] = (
                    attention_mask
                )


        decoder_outputs = decoder_session.run(
            None,
            decoder_inputs
        )

        logits = decoder_outputs[0]

        next_token_logits = (
            logits[0, -1, :]
        )

        next_token = int(
            np.argmax(
                next_token_logits
            )
        )

        generated.append(
            next_token
        )

        if next_token == EOS_ID:
            break


    # -------------------------
    # DECODE OUTPUT
    # -------------------------

    output_ids = generated[1:]

    if EOS_ID in output_ids:

        output_ids = output_ids[
            :output_ids.index(EOS_ID)
        ]

    result = sp.decode(
        output_ids
    )

    return result


# =========================
# GRADIO UI
# =========================

with gr.Blocks(
    title="AnViMt - English to Vietnamese Translator"
) as demo:

    gr.Markdown(
        """
# 🇻🇳 AnViMt

### English → Vietnamese Machine Translation

Nhập văn bản tiếng Anh và AnViMt sẽ dịch sang tiếng Việt.

**Model:** Phitran21/AnViMt  
**Architecture:** BART Seq2Seq  
**Inference:** ONNX Runtime
"""
    )

    with gr.Row():

        input_text = gr.Textbox(
            label="English",
            placeholder=(
                "Enter English text..."
            ),
            lines=8
        )

        output_text = gr.Textbox(
            label="Vietnamese",
            lines=8
        )


    translate_button = gr.Button(
        "Translate",
        variant="primary"
    )


    translate_button.click(
        fn=translate,
        inputs=input_text,
        outputs=output_text
    )


    input_text.submit(
        fn=translate,
        inputs=input_text,
        outputs=output_text
    )


    gr.Examples(
        examples=[
            [
                "The system automatically restarts the service after an unexpected process crash."
            ],
            [
                "A lone swordsman entered the abandoned temple at midnight, unaware that his greatest enemy was waiting in the shadows."
            ],
            [
                "Could you please restart the service and check whether the error still occurs?"
            ],
        ],
        inputs=input_text
    )


demo.launch()