# app.py — HF Spaces Free (CPU), Hunyuan-MT 7B-fp8, đa ngôn ngữ, chia đoạn, UI + API
import os, re
from typing import List, Optional

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ===== Cấu hình =====
DEFAULT_MODEL = "tencent/Hunyuan-MT-7B-fp8"  # đổi bằng env MODEL_NAME nếu muốn
MODEL_NAME = os.getenv("MODEL_NAME", DEFAULT_MODEL)

GEN_KW = dict(  # tham số sinh nhẹ cho CPU
    max_new_tokens=256,
    top_k=20,
    top_p=0.6,
    repetition_penalty=1.05,
    temperature=0.7,
    do_sample=True,
)

MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "800"))  # giới hạn input mỗi mảnh

# ===== Load tokenizer & model (fp8 bằng dict quantization_config) =====
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
quant_cfg = {"quantization_method": "fp8", "ignore": []}  # tránh lỗi ignore=None
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    quantization_config=quant_cfg,
)
DEVICE = getattr(model, "device", torch.device("cpu"))

# ===== Chuẩn hóa tên ngôn ngữ =====
LANG_ALIASES = {
    "vi": "Vietnamese", "vie": "Vietnamese", "vietnamese": "Vietnamese", "tiếng việt": "Vietnamese",
    "zh": "Chinese", "chi": "Chinese", "zho": "Chinese", "chinese": "Chinese", "tiếng trung": "Chinese", "hán ngữ": "Chinese", "mandarin": "Chinese",
    "en": "English", "eng": "English", "tiếng anh": "English", "english": "English",
    "ja": "Japanese", "jpn": "Japanese", "tiếng nhật": "Japanese", "japanese": "Japanese",
    "ko": "Korean", "kor": "Korean", "tiếng hàn": "Korean", "korean": "Korean",
    "fr": "French", "fra": "French", "fre": "French", "tiếng pháp": "French", "french": "French",
    "de": "German", "deu": "German", "ger": "German", "tiếng đức": "German", "german": "German",
    "es": "Spanish", "spa": "Spanish", "tiếng tây ban nha": "Spanish", "spanish": "Spanish",
    "th": "Thai", "tha": "Thai", "tiếng thái": "Thai", "thai": "Thai",
    "id": "Indonesian", "ind": "Indonesian", "tiếng indonesia": "Indonesian", "indonesian": "Indonesian",
    "ms": "Malay", "msa": "Malay", "tiếng malaysia": "Malay", "malay": "Malay",
    "pt": "Portuguese", "por": "Portuguese", "tiếng bồ đào nha": "Portuguese", "portuguese": "Portuguese",
    "ru": "Russian", "rus": "Russian", "tiếng nga": "Russian", "russian": "Russian",
}
LANG_CHOICES = sorted(set(LANG_ALIASES.values()))
def norm_lang(s: Optional[str]) -> Optional[str]:
    if not s: return None
    k = s.strip().lower()
    return LANG_ALIASES.get(k, s.strip())

# ===== Chia văn bản theo token =====
def chunk_by_tokens(text: str, max_tokens: int) -> List[str]:
    text = text.strip()
    if not text: return []
    rough = re.split(r"(?<=[\.!?。！？])\s+", text)
    chunks, buf = [], ""
    def tok_len(s: str) -> int:
        return tokenizer(s, add_special_tokens=False, return_length=True)["length"]
    for part in rough:
        cand = (buf + " " + part).strip() if buf else part
        if tok_len(cand) <= max_tokens:
            buf = cand
        else:
            if buf: chunks.append(buf); buf = ""
            if tok_len(part) <= max_tokens:
                buf = part
            else:
                ids = tokenizer(part, add_special_tokens=False)["input_ids"]
                for i in range(0, len(ids), max_tokens):
                    piece = tokenizer.decode(ids[i:i+max_tokens], skip_special_tokens=True)
                    if piece.strip(): chunks.append(piece.strip())
    if buf: chunks.append(buf)
    return [c for c in chunks if c.strip()]

# ===== Core translate (chat template) =====
@torch.inference_mode()
def translate_text(text: str, target_lang: str, source_lang: Optional[str]=None) -> str:
    tgt = norm_lang(target_lang) or "Vietnamese"
    src = norm_lang(source_lang)
    sys_prompt = (f"Translate the following segment from {src} into {tgt}, without additional explanation."
                  if src else
                  f"Translate the following segment into {tgt}, without additional explanation.")
    outs = []
    for piece in chunk_by_tokens(text, MAX_INPUT_TOKENS):
        msgs = [{"role":"user","content": f"{sys_prompt}\n\n{piece}"}]
        inputs = tokenizer.apply_chat_template(msgs, tokenize=True, add_generation_prompt=False, return_tensors="pt")
        out_ids = model.generate(inputs.to(DEVICE), **GEN_KW)
        outs.append(tokenizer.decode(out_ids[0], skip_special_tokens=True).strip())
    return "\n".join(outs).strip()

def translate_batch(texts: List[str], target_lang: str, source_lang: Optional[str]=None) -> List[str]:
    return [translate_text(t, target_lang, source_lang) for t in texts]

# ===== Gradio UI + API =====
with gr.Blocks() as demo:
    gr.Markdown("## Hunyuan-MT 7B-fp8 — Multilingual Translation (HF Free CPU)\nChia đoạn theo token, UI + API (Gradio).")

    with gr.Tab("Single"):
        src = gr.Textbox(label="Văn bản nguồn", lines=10, placeholder="Dán văn bản cần dịch…")
        with gr.Row():
            src_lang = gr.Textbox(label="Ngôn ngữ nguồn (tùy chọn)", placeholder="Ví dụ: Vietnamese/Chinese/English…")
            tgt_lang = gr.Dropdown(label="Ngôn ngữ đích", choices=LANG_CHOICES, value="Vietnamese")
        out = gr.Textbox(label="Bản dịch", lines=10)
        gr.Button("Dịch").click(translate_text, inputs=[src, tgt_lang, src_lang], outputs=out, api_name="translate_text")

    with gr.Tab("Batch"):
        src_list = gr.Textbox(label="Mỗi dòng 1 câu/đoạn", lines=10)
        with gr.Row():
            src_lang_b = gr.Textbox(label="Ngôn ngữ nguồn (tùy chọn)")
            tgt_lang_b = gr.Dropdown(label="Ngôn ngữ đích", choices=LANG_CHOICES, value="Vietnamese")
        out_list = gr.Textbox(label="Kết quả (mỗi dòng tương ứng 1 đầu vào)", lines=10)
        def _batch(txts_raw: str, tgt: str, src_: Optional[str]):
            texts = [x for x in txts_raw.splitlines() if x.strip()]
            return "\n".join(translate_batch(texts, tgt, src_))
        gr.Button("Dịch Batch").click(_batch, inputs=[src_list, tgt_lang_b, src_lang_b], outputs=out_list, api_name="translate_batch")

demo.queue(concurrency_count=1, max_size=2).launch()