"""
AI Translation System với Knowledge Graph + NER + Transliteration
==================================================================
Pipeline 4 lớp xử lý tên riêng:
  Lớp 1: Lookup Knowledge Graph DB (entity đã biết, độ chính xác cao nhất)
  Lớp 2: NER tự động phát hiện tên chưa có trong DB
  Lớp 3: Transliteration (unidecode) cho PERSON/LOCATION
  Lớp 4: NMT Model dịch phần còn lại (tên đã được mask)
"""

import gradio as gr
import os
import re
import logging
from dataclasses import dataclass, field
from functools import lru_cache

import httpx
from unidecode import unidecode

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("translation_app")


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
@dataclass
class AppConfig:
    supabase_url: str = field(
        default_factory=lambda: os.environ.get(
            "SUPABASE_URL",
            "https://jkaazcjahfqozwizhpvk.supabase.co"
        )
    )
    supabase_key: str = field(
        default_factory=lambda: os.environ.get("SUPABASE_KEY", "")
    )
    allowed_types: tuple = (
        "PERSON", "ORG", "LOCATION",
        "TECH", "PROGRAMMING_LANGUAGE", "FRAMEWORK",
        "DATABASE", "CLOUD", "MODEL",
        "PRODUCT", "FOOD",
    )
    min_frequency: int = 1          # Hạ xuống 1 để không bỏ sót entity
    
    # --- ĐÃ CẬP NHẬT MODEL CỦA BẠN TẠI ĐÂY ---
    ev_model: str = "LLC536/nmt-en-vi-translation-v1"
    ve_model: str = "LLC536/nmt-vi-en-translation-v1"
    # -----------------------------------------

CONFIG = AppConfig()


# ---------------------------------------------------------------------------
# Model – Lazy Loading
# ---------------------------------------------------------------------------
_models: dict = {}


def get_model(direction: str):
    if direction not in _models:
        from transformers import pipeline
        model_id = CONFIG.ev_model if direction == "en->vi" else CONFIG.ve_model
        logger.info(f"Đang load model dịch: {model_id}")
        _models[direction] = pipeline("translation", model=model_id)
        logger.info(f"Model [{direction}] sẵn sàng.")
    return _models[direction]


# ---------------------------------------------------------------------------
# NER – Lazy Loading (tránh import nặng lúc startup)
# ---------------------------------------------------------------------------
_ner_vi = None
_ner_en = None


def get_ner_vi():
    """NER tiếng Việt dùng underthesea."""
    global _ner_vi
    if _ner_vi is None:
        try:
            from underthesea import ner as _ner
            _ner_vi = _ner
            logger.info("NER tiếng Việt (underthesea) sẵn sàng.")
        except ImportError:
            logger.warning("underthesea chưa cài. NER tiếng Việt bị tắt.")
            _ner_vi = None
    return _ner_vi


def get_ner_en():
    """NER tiếng Anh dùng transformers pipeline."""
    global _ner_en
    if _ner_en is None:
        try:
            from transformers import pipeline
            _ner_en = pipeline(
                "ner",
                model="dslim/bert-base-NER",
                aggregation_strategy="simple",
            )
            logger.info("NER tiếng Anh (dslim/bert-base-NER) sẵn sàng.")
        except Exception as e:
            logger.warning(f"Không load được NER tiếng Anh: {e}")
            _ner_en = None
    return _ner_en


# ---------------------------------------------------------------------------
# Supabase REST API
# ---------------------------------------------------------------------------
@lru_cache(maxsize=8)
def fetch_entities_from_supabase(lang_code: str) -> dict:
    """
    Lấy entity glossary từ Supabase theo ngôn ngữ.
    Cache lru_cache – chỉ gọi API 1 lần mỗi lang_code.
    """
    if not CONFIG.supabase_key:
        logger.warning("SUPABASE_KEY chưa cấu hình. KG bị bỏ qua.")
        return {}

    headers = {
        "apikey": CONFIG.supabase_key,
        "Authorization": f"Bearer {CONFIG.supabase_key}",
    }
    params = {
        "select": "entity_name,normalize_name",
        "language": f"eq.{lang_code}",
        "frequency": f"gte.{CONFIG.min_frequency}",
        "entity_type": f"in.({','.join(CONFIG.allowed_types)})",
    }

    try:
        resp = httpx.get(
            f"{CONFIG.supabase_url}/rest/v1/knowledge_graph_entities",
            headers=headers,
            params=params,
            timeout=10.0,
        )
        resp.raise_for_status()
        entities = {}
        for row in resp.json():
            name = (row.get("entity_name") or "").strip()
            norm = (row.get("normalize_name") or "").strip()
            if name and norm:
                entities[name] = norm
        logger.info(f"Loaded {len(entities)} entities từ Supabase (lang={lang_code}).")
        return entities
    except httpx.HTTPStatusError as e:
        logger.error(f"Supabase HTTP {e.response.status_code}: {e.response.text}")
        return {}
    except Exception as e:
        logger.error(f"Lỗi Supabase API: {e}", exc_info=True)
        return {}


def invalidate_cache():
    fetch_entities_from_supabase.cache_clear()
    logger.info("Entity cache đã xóa.")


# ---------------------------------------------------------------------------
# Lớp 2+3: NER tự động + Transliteration
# ---------------------------------------------------------------------------
def _convert_by_type(name: str, entity_type: str) -> str:
    """
    Quyết định cách xử lý tên theo loại thực thể:
    - PERSON / LOC  → bỏ dấu bằng unidecode
    - ORG / TECH    → giữ nguyên (tên tổ chức/kỹ thuật không dịch)
    """
    if entity_type in ("PER", "PERSON", "LOC", "LOCATION"):
        return unidecode(name)   # "Lê Hoàng Sơn" → "Le Hoang Son"
    return name                  # "Google", "Python" → giữ nguyên


def auto_detect_entities_vi(text: str) -> dict:
    """
    Dùng underthesea NER để phát hiện tên riêng trong tiếng Việt.
    Trả về {tên_gốc: tên_đích}.
    """
    ner_fn = get_ner_vi()
    if not ner_fn:
        return {}

    try:
        results = ner_fn(text)
        # underthesea trả về list of (word, pos, chunk, ner_tag)
        # hoặc list of (word, ner_tag) tuỳ version
        detected = {}
        current_words = []
        current_type  = None

        for item in results:
            # Hỗ trợ cả 2 format output của underthesea
            word = item[0]
            tag  = item[-1]  # tag luôn ở cuối

            if tag.startswith("B-"):
                # Lưu entity trước đó
                if current_words and current_type:
                    name = " ".join(current_words)
                    detected[name] = _convert_by_type(name, current_type)
                current_words = [word]
                current_type  = tag[2:]   # "PER", "LOC", "ORG"...

            elif tag.startswith("I-") and current_words:
                current_words.append(word)

            else:
                # Token không phải entity → flush
                if current_words and current_type:
                    name = " ".join(current_words)
                    detected[name] = _convert_by_type(name, current_type)
                current_words = []
                current_type  = None

        # Flush entity cuối cùng
        if current_words and current_type:
            name = " ".join(current_words)
            detected[name] = _convert_by_type(name, current_type)

        logger.info(f"NER VI phát hiện {len(detected)} entity: {list(detected.keys())}")
        return detected

    except Exception as e:
        logger.error(f"Lỗi NER tiếng Việt: {e}", exc_info=True)
        return {}


def auto_detect_entities_en(text: str) -> dict:
    """
    Dùng bert-base-NER để phát hiện tên riêng trong tiếng Anh.
    Trả về {tên_gốc: tên_đích} (giữ nguyên – EN→VI thì model lo).
    """
    ner_fn = get_ner_en()
    if not ner_fn:
        return {}

    try:
        results = ner_fn(text)
        detected = {}
        for ent in results:
            word      = ent.get("word", "").strip()
            ent_type  = ent.get("entity_group", "")
            score     = ent.get("score", 0)
            if word and score > 0.85:
                # Giữ nguyên tên tiếng Anh khi dịch sang tiếng Việt
                detected[word] = word
        logger.info(f"NER EN phát hiện {len(detected)} entity: {list(detected.keys())}")
        return detected
    except Exception as e:
        logger.error(f"Lỗi NER tiếng Anh: {e}", exc_info=True)
        return {}


def auto_detect_entities(text: str, src_lang: str) -> dict:
    """Dispatcher: chọn NER engine theo ngôn ngữ nguồn."""
    if src_lang == "vi":
        return auto_detect_entities_vi(text)
    return auto_detect_entities_en(text)


# ---------------------------------------------------------------------------
# Mask / Restore entity
# ---------------------------------------------------------------------------
def mask_entities(text: str, entity_dict: dict) -> tuple:
    """
    Thay tên riêng bằng placeholder trước khi đưa vào model dịch.
    Trả về (text_đã_mask, {placeholder: tên_đích})

    Ví dụ:
      Input:  "Lê Hoàng Sơn là kỹ sư"
      Output: ("__E0__ là kỹ sư", {"__E0__": "Le Hoang Son"})
    """
    placeholder_map = {}
    masked = text
    sorted_names = sorted(entity_dict.keys(), key=len, reverse=True)

    for i, name in enumerate(sorted_names):
        # So sánh case-insensitive khi tìm kiếm
        pattern = re.compile(re.escape(name), re.IGNORECASE | re.UNICODE)
        if pattern.search(masked):
            token = f"__E{i}__"
            placeholder_map[token] = entity_dict[name]
            masked = pattern.sub(token, masked)

    return masked, placeholder_map


def restore_entities(text: str, placeholder_map: dict) -> str:
    """Thay placeholder ngược lại về tên đúng sau khi dịch."""
    for token, correct_name in placeholder_map.items():
        text = text.replace(token, correct_name)
    return text


# ---------------------------------------------------------------------------
# Smart Replace (post-processing fallback)
# ---------------------------------------------------------------------------
def smart_replace(text: str, entity_dict: dict) -> str:
    """
    Post-processing: replace entity còn sót trong output
    (những tên model tự dịch dù không qua mask).
    """
    if not entity_dict:
        return text
    for ent in sorted(entity_dict, key=len, reverse=True):
        pattern = re.compile(
            rf"(?<!\w){re.escape(ent)}(?!\w)",
            re.IGNORECASE | re.UNICODE,
        )
        text = pattern.sub(entity_dict[ent], text)
    return text


# ---------------------------------------------------------------------------
# Core Translation Logic – Pipeline 4 lớp
# ---------------------------------------------------------------------------
def translate_logic(text: str, mode: str, use_kg: bool) -> str:
    if not text or not text.strip():
        return ""

    direction   = "en->vi" if mode == "Anh -> Việt" else "vi->en"
    src_lang    = "en"     if direction == "en->vi" else "vi"
    target_lang = "vi"     if direction == "en->vi" else "en"

    if not use_kg:
        try:
            return get_model(direction)(text)[0]["translation_text"]
        except Exception as e:
            return f"[Lỗi model]: {e}"

    try:
        # ── Lớp 1: Knowledge Graph DB ─────────────────────────────────────
        # Entity đã biết, độ chính xác cao nhất
        db_entities = fetch_entities_from_supabase(src_lang)
        db_match    = {k: v for k, v in db_entities.items()
                       if re.search(re.escape(k), text, re.IGNORECASE)}

        # ── Lớp 2+3: NER tự động + Transliteration ────────────────────────
        # Phát hiện tên chưa có trong DB, rồi bỏ dấu/giữ nguyên theo loại
        ner_entities = auto_detect_entities(text, src_lang)

        # Merge: DB ưu tiên hơn NER nếu cùng key
        all_entities = {**ner_entities, **db_match}
        logger.info(
            f"Tổng entity sẽ mask: {len(all_entities)} "
            f"(DB={len(db_match)}, NER={len(ner_entities)})"
        )

        # ── Lớp 4: NMT Model ──────────────────────────────────────────────
        # Mask tên → dịch → restore
        masked_text, placeholder_map = mask_entities(text, all_entities)
        logger.info(f"Text sau mask: {masked_text}")

        raw_output = get_model(direction)(masked_text)[0]["translation_text"]
        final_output = restore_entities(raw_output, placeholder_map)

        # Post-processing: catch entity đích còn sót (defense in depth)
        tgt_entities = fetch_entities_from_supabase(target_lang)
        tgt_match    = {k: v for k, v in tgt_entities.items() if k in final_output}
        final_output = smart_replace(final_output, tgt_match)

        return final_output

    except Exception as e:
        logger.error(f"Lỗi pipeline: {e}", exc_info=True)
        try:
            return get_model(direction)(text)[0]["translation_text"]
        except Exception as e2:
            return f"[Lỗi model]: {e2}"


# ---------------------------------------------------------------------------
# Gradio UI (tương thích Gradio 5.x)
# ---------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="AI Translation + KG") as demo:
    gr.Markdown("# 🚀 Hệ thống Dịch thuật AI tích hợp Knowledge Graph + NER")
    gr.Markdown(
        "Dịch Anh ↔ Việt với pipeline 4 lớp bảo vệ thực thể:\n"
        "**KG DB** → **NER tự động** → **Transliteration** → **NMT Model**"
    )

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                lines=7,
                label="Văn bản gốc",
                placeholder="Nhập câu cần dịch tại đây…",
            )
            mode = gr.Radio(
                choices=["Anh -> Việt", "Việt -> Anh"],
                label="Chiều dịch",
                value="Anh -> Việt",
            )
            use_kg = gr.Checkbox(
                label="Bật Knowledge Graph + NER + Transliteration",
                value=True,
            )
            btn = gr.Button("🔄 Dịch ngay", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(
                lines=7,
                label="Kết quả dịch",
                interactive=False,
            )
            gr.Markdown(
                "**Pipeline khi bật KG:**\n\n"
                "1. 🗄️ Tra Knowledge Graph DB\n"
                "2. 🔍 NER tự động phát hiện tên mới\n"
                "3. 🔤 Transliteration: bỏ dấu tên riêng\n"
                "4. 🤖 Model dịch (tên đã được bảo vệ)\n\n"
                "→ Không còn lỗi `Lê Hoàng Sơn` → `Li Wongshan` ✅"
            )

    with gr.Accordion("⚙️ Quản lý cache (admin)", open=False):
        refresh_btn = gr.Button("🗑️ Xóa entity cache (reload từ Supabase)")
        cache_status = gr.Textbox(label="Trạng thái", interactive=False)

        def do_refresh():
            invalidate_cache()
            return " Cache đã xóa. Entity sẽ reload ở request tiếp theo."

        refresh_btn.click(fn=do_refresh, outputs=cache_status)

    btn.click(
        fn=translate_logic,
        inputs=[input_text, mode, use_kg],
        outputs=output_text,
    )

if __name__ == "__main__":
    demo.launch()