import os
import time
from collections import defaultdict, deque
from typing import Tuple

import gradio as gr
from huggingface_hub import InferenceClient

# =========================
# CONFIG & AUTH
# =========================
# Lấy token từ biến môi trường (khuyên dùng: add tại Settings → Repository secrets)
# Nếu cần hardcode để test nhanh (KHÔNG an toàn repo public):
# HF_TOKEN = "hf_XXXXXXXXXXXXXXXXXXXXXXXX"  # <--- DÁN TOKEN Ở ĐÂY (không khuyến nghị)
HF_TOKEN = os.getenv("HF_TOKEN", None)

# Model mặc định cho UI (có thể đổi qua dropdown)
DEFAULT_MODEL = os.getenv("HYMT_MODEL", "tencent/Hunyuan-MT-7B-fp8")

# Model cố định cho API
FIXED_MODEL = "tencent/Hunyuan-MT-7B-fp8"

# Giới hạn đồng thời theo event (thay cho concurrency_count cũ)
UI_CONCURRENCY_LIMIT = int(os.getenv("UI_CONCURRENCY_LIMIT", "2"))   # nút Dịch
API_CONCURRENCY_LIMIT = int(os.getenv("API_CONCURRENCY_LIMIT", "2")) # endpoint API

# Số worker tổng (tuỳ chọn): ảnh hưởng thread pool của server
LAUNCH_MAX_THREADS = int(os.getenv("LAUNCH_MAX_THREADS", "40"))

# =========================
# NGÔN NGỮ & PROMPT
# =========================
LANGS = [
    ("Chinese (简体中文)", "zh"),
    ("Traditional Chinese (繁體中文)", "zh-Hant"),
    ("Cantonese (粤语)", "yue"),
    ("English (English)", "en"),
    ("Vietnamese (Tiếng Việt)", "vi"),
    ("Japanese (日本語)", "ja"),
    ("Korean (한국어)", "ko"),
    ("Thai (ไทย)", "th"),
    ("French (Français)", "fr"),
    ("Spanish (Español)", "es"),
    ("Portuguese (Português)", "pt"),
    ("Italian (Italiano)", "it"),
    ("German (Deutsch)", "de"),
    ("Russian (Русский)", "ru"),
    ("Arabic (العربية)", "ar"),
    ("Turkish (Türkçe)", "tr"),
    ("Indonesian (Bahasa Indonesia)", "id"),
    ("Malay (Bahasa Melayu)", "ms"),
    ("Filipino (Filipino)", "tl"),
    ("Hindi (हिन्दी)", "hi"),
    ("Polish (Polski)", "pl"),
    ("Czech (Čeština)", "cs"),
    ("Dutch (Nederlands)", "nl"),
    ("Khmer (ភាសាខ្មែរ)", "km"),
    ("Burmese (မြန်မာ)", "my"),
    ("Persian (فارسی)", "fa"),
    ("Gujarati (ગુજરાતી)", "gu"),
    ("Urdu (اردو)", "ur"),
    ("Telugu (తెలుగు)", "te"),
    ("Marathi (मराठी)", "mr"),
    ("Hebrew (עברית)", "he"),
    ("Bengali (বাংলা)", "bn"),
    ("Tamil (தமிழ்)", "ta"),
    ("Ukrainian (Українська)", "uk"),
    ("Tibetan (བོད་ཡིག)", "bo"),
    ("Kazakh (Қазақша)", "kk"),
    ("Mongolian (Монгол)", "mn"),
    ("Uyghur (ئۇيغۇرچە)", "ug"),
]
LABEL2CODE = {label: code for label, code in LANGS}
ZH_CODES = {"zh", "zh-Hant", "yue"}


def build_prompt(src_lang: str, tgt_lang: str, text: str) -> str:
    """
    Prompt template:
    - Nếu có tiếng Trung (zh/zh-Hant/yue) ở nguồn hoặc đích -> template tiếng Trung
    - Nếu không -> template tiếng Anh
    """
    txt = (text or "").strip()
    if not txt:
        return ""
    if src_lang in ZH_CODES or tgt_lang in ZH_CODES:
        return f"把下面的文本翻译成{tgt_lang}，不要额外解释。\n\n{txt}"
    else:
        return f"Translate the following segment into {tgt_lang}, without additional explanation.\n\n{txt}"


# =========================
# INFERENCE HELPER
# =========================
def call_hf_inference(model: str, prompt: str) -> str:
    """
    Gọi Hugging Face Serverless Inference API (text-generation).
    Không cần GPU, phù hợp Space free.
    """
    if not prompt:
        return ""
    client = InferenceClient(token=HF_TOKEN)
    try:
        out = client.text_generation(
            model=model,
            prompt=prompt,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.6,
            repetition_penalty=1.05,
            stream=False,
        )
        return (out or "").strip()
    except Exception as e:
        return f"[Lỗi] Không thể gọi Inference API: {e}"


def translate(text: str, src_code: str, tgt_code: str, model_choice: str) -> str:
    if not text or not text.strip():
        return "Vui lòng nhập nội dung cần dịch."
    if not src_code or not tgt_code:
        return "Thiếu mã ngôn ngữ nguồn/đích."
    if src_code == tgt_code:
        return text.strip()
    prompt = build_prompt(src_code, tgt_code, text)
    return call_hf_inference(model_choice, prompt)


# =========================
# RATE LIMIT THEO IP (IN-MEMORY)
# =========================
RATE_WINDOW_SEC = int(os.getenv("RATE_WINDOW_SEC", "60"))  # ví dụ: 60 giây
RATE_MAX_REQ = int(os.getenv("RATE_MAX_REQ", "10"))        # ví dụ: 10 request / IP / 60s

_ip_buckets: dict[str, deque] = defaultdict(deque)

def _rate_limited(request: gr.Request) -> Tuple[bool, str]:
    """
    Trả (ok, msg). ok=False nếu vượt ngưỡng.
    Lưu ý: in-memory -> reset khi container restart.
    """
    try:
        ip = (request.client.host if request and request.client else "unknown") or "unknown"
    except Exception:
        ip = "unknown"

    now = time.time()
    dq = _ip_buckets[ip]

    # Bỏ các dấu thời gian ngoài cửa sổ
    while dq and (now - dq[0] > RATE_WINDOW_SEC):
        dq.popleft()

    if len(dq) >= RATE_MAX_REQ:
        wait_sec = max(1, int(RATE_WINDOW_SEC - (now - dq[0])))
        return False, f"Bạn đã vượt giới hạn {RATE_MAX_REQ} yêu cầu / {RATE_WINDOW_SEC}s. Hãy thử lại sau ~{wait_sec}s."
    dq.append(now)
    return True, ""


# =========================
# UI (GRADIO BLOCKS)
# =========================
def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Hunyuan-MT Translation (HF Inference API)", fill_height=True) as demo:
        gr.Markdown(
            """
            # Tencent Hunyuan-MT (Serverless)
            Chạy trên **Hugging Face Space (CPU free)** bằng **Serverless Inference API**.
            - Chọn mô hình `tencent/Hunyuan-MT-7B-fp8` (khuyến nghị) hoặc `tencent/Hunyuan-MT-7B`.
            - Chọn ngôn ngữ nguồn/đích rồi bấm **Dịch**.
            > Mẹo: vào *Settings → Repository secrets* thêm `HF_TOKEN` để tăng hạn mức.
            """
        )

        with gr.Row():
            model_choice = gr.Dropdown(
                choices=[
                    "tencent/Hunyuan-MT-7B-fp8",
                    "tencent/Hunyuan-MT-7B",
                ],
                value=DEFAULT_MODEL,
                label="Model (Serverless)"
            )

        with gr.Row():
            src = gr.Dropdown(
                choices=[l for l, _ in LANGS],
                value="English (English)",
                label="Nguồn"
            )
            tgt = gr.Dropdown(
                choices=[l for l, _ in LANGS],
                value="Vietnamese (Tiếng Việt)",
                label="Đích"
            )

        inp = gr.Textbox(label="Nội dung cần dịch", lines=8, placeholder="Nhập văn bản…")
        out = gr.Textbox(label="Kết quả", lines=8)
        btn = gr.Button("Dịch", variant="primary")

        def _on_translate(text, src_label, tgt_label, model_id, request: gr.Request = None):
            ok, msg = _rate_limited(request)
            if not ok:
                return msg
            src_code = LABEL2CODE[src_label]
            tgt_code = LABEL2CODE[tgt_label]
            return translate(text, src_code, tgt_code, model_id)

        # ✅ Đặt concurrency_limit ngay trên event listener (chuẩn mới)
        btn.click(
            _on_translate,
            [inp, src, tgt, model_choice],
            [out],
            concurrency_limit=UI_CONCURRENCY_LIMIT
        )

        gr.Markdown(
            """
            #### Lưu ý
            - Demo gọi **Serverless Inference API** nên tốc độ phụ thuộc hạn mức.
            - Cần throughput cao hơn? Hãy cân nhắc GPU hoặc tự triển khai TGI/vLLM.
            """
        )
    return demo


# =========================
# API CỐ ĐỊNH (MODEL KHÓA)
# =========================
def api_translate_fixed(text: str, src_code: str, tgt_code: str, request: gr.Request = None) -> str:
    """
    API cho website: nhận mã ngôn ngữ (vd: 'en', 'vi', 'zh', ...),
    dùng model cố định Hunyuan-MT-7B-fp8.
    """
    ok, msg = _rate_limited(request)
    if not ok:
        return msg
    if not text or not text.strip():
        return ""
    if not src_code or not tgt_code:
        return "Thiếu mã ngôn ngữ nguồn/đích."
    if src_code == tgt_code:
        return text.strip()
    prompt = build_prompt(src_code, tgt_code, text)
    return call_hf_inference(FIXED_MODEL, prompt)


def build_api_interface() -> gr.Interface:
    """
    Interface riêng để có endpoint REST.
    - /run/api_translate_fixed (nếu platform hỗ trợ)
    - Hoặc /run/predict với fn_index tương ứng
    """
    return gr.Interface(
        fn=api_translate_fixed,
        inputs=[
            gr.Textbox(label="text"),
            gr.Textbox(label="src_code"),
            gr.Textbox(label="tgt_code"),
        ],
        outputs=gr.Textbox(label="translation"),
        title="Hunyuan-MT Fixed API",
        description="POST JSON tới endpoint để nhận bản dịch. Model cố định: tencent/Hunyuan-MT-7B-fp8.",
        concurrency_limit=API_CONCURRENCY_LIMIT,  # ✅ giới hạn đồng thời cho API
    )


# =========================
# MAIN
# =========================
if __name__ == "__main__":
    ui_app = build_ui()
    api_iface = build_api_interface()

    # Gộp UI + API vào cùng server
    demo = gr.TabbedInterface([ui_app, api_iface], tab_names=["App", "API"])

    # Hàng chờ: KHÔNG dùng concurrency_count nữa!
    QUEUE_MAX = int(os.getenv("GRADIO_QUEUE_MAX", "20"))  # số job có thể chờ
    demo = demo.queue(max_size=QUEUE_MAX, status_update_rate=2)

    # Bật REST API; (tuỳ chọn) khống chế thread tổng bằng max_threads
    demo.launch(enable_api=True, max_threads=LAUNCH_MAX_THREADS)