import os
import numpy as np
import onnxruntime as ort
import sentencepiece as spm


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

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

ENCODER_PATH = os.path.join(BASE_DIR, "encoder_model.onnx")
DECODER_PATH = os.path.join(BASE_DIR, "decoder_model.onnx")
SPM_PATH = os.path.join(BASE_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: vocab={sp.get_piece_size()} | "
    f"PAD={sp.pad_id()} BOS={sp.bos_id()} EOS={sp.eos_id()}"
)


# =========================
# ONNX RUNTIME
# =========================

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

providers = ["CPUExecutionProvider"]

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

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


# =========================
# SHOW MODEL INFO
# =========================

print("\nEncoder inputs:")
for x in encoder_session.get_inputs():
    print(" ", x.name, x.shape, x.type)

print("\nDecoder inputs:")
for x in decoder_session.get_inputs():
    print(" ", x.name, x.shape, x.type)


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

def translate(text):
    text = text.strip()

    if not text:
        return ""

    # ---------------------------------
    # 1. ENCODE INPUT
    # ---------------------------------

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

    # Giới hạn context
    src_ids = src_ids[:MAX_INPUT_LEN - 2]

    # BART format
    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
    )

    # ---------------------------------
    # 2. ENCODER
    # ---------------------------------

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

    encoder_hidden_states = encoder_outputs[0]

    # ---------------------------------
    # 3. DECODER GREEDY
    # ---------------------------------

    generated = [BOS_ID]

    for _ in range(MAX_OUTPUT_LEN):

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

        decoder_inputs = {}

        # Tự map input để tránh lệch tên
        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
        )

        # Output đầu tiên là logits
        logits = decoder_outputs[0]

        # Lấy token cuối
        next_token_logits = logits[0, -1, :]

        next_token = int(
            np.argmax(next_token_logits)
        )

        generated.append(next_token)

        if next_token == EOS_ID:
            break

    # ---------------------------------
    # 4. DECODE
    # ---------------------------------

    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


# =========================
# INTERACTIVE CLI
# =========================

print("\nEN → VI translator v2 ready")
print("Type 'exit' để thoát\n")

while True:

    try:
        text = input("EN > ").strip()

    except (KeyboardInterrupt, EOFError):
        print("\nBye")
        break

    if text.lower() in ["exit", "quit", "q"]:
        print("Bye")
        break

    if not text:
        continue

    result = translate(text)

    print("VI >", result)
    print()