"""번역(MT) 래퍼.

두 백엔드를 지원한다.

* CTranslate2 — NLLB 계열 seq2seq. 가볍고 빠르며 추가학습이 쉽다.
* Transformers — Seed-X / Qwen3 같은 LLM. 품질이 좋고 LoRA 어댑터를 얹을 수 있다.

LLM 백엔드는 다시 두 가지 프롬프트 형식으로 갈린다. 같은 "LLM"이라도 요구사항이
정반대라 하나로 합치면 한쪽이 반드시 망가진다.

* PromptStyle.SEEDX — Seed-X 는 chat template 이 없는 **번역 전용 completion 모델**이다.
  모델 카드가 정한 고정 문장과 끝의 `<언어코드>` 태그를 그대로 지켜야 하고,
  그 밖의 지시문을 끼워 넣으면 학습 분포를 벗어난다. 따라서 용어집은 프롬프트가
  아니라 **플레이스홀더 치환**으로 넣는다.
* PromptStyle.INSTRUCT — Qwen3 등. chat template 을 쓰고 용어집을 지시문으로 넣는다.

세 경우 모두 `Translator` 인터페이스를 따르므로 상위 레이어는 구분하지 않는다.
"""

from __future__ import annotations

import abc
import logging
import re
import threading

from ..constants import LANGUAGES, models_dir
from .glossary import Glossary
from .speech_level import Politeness, SpeechLevel, fix_particles, prompt_instruction
from .tiers import MTBackend, MTSpec, PromptStyle

log = logging.getLogger(__name__)

#: 문장 끝에 LLM이 붙이곤 하는 군더더기를 걷어내기 위한 패턴
_TRAILING_NOISE = re.compile(r"^\s*(?:번역[:：]|Translation[:：])\s*", re.IGNORECASE)


def _restore_or_retranslate(out: str, replacements: list[str], retry) -> str:
    """자리표시자를 역어로 되돌린다. 하나라도 사라졌으면 용어집 없이 다시 번역.

    번역 모델이 문장 일부를 통째로 누락하면 그 안의 자리표시자도 같이
    사라진다. 그대로 복원하면 용어가 자막에서 증발하는데, 그건 용어집을
    아예 안 쓴 것보다 나쁘다. 그래서 그럴 때는 보호 없이 한 번 더 번역해
    원문 내용이 살아 있는 결과를 쓴다.
    """
    missing = Glossary.missing_placeholders(out, replacements)
    if not missing:
        return Glossary.restore(out, replacements)
    log.debug("자리표시자 %s 소실 — 용어집 없이 재번역", missing)
    return retry()


def build_seedx_prompt(text: str, source_lang: str, target_lang: str) -> str:
    """Seed-X 모델 카드가 요구하는 정확한 프롬프트.

        Translate the following English sentence into Chinese:
        May the force be with you <zh>

    끝의 `<언어코드>` 태그는 PPO 학습에 쓰인 것이라 **필수**다. 빠지면 번역
    품질이 크게 흔들린다. 형식을 바꾸면 안 되므로 지시문이나 용어집 설명을
    여기 끼워 넣지 말 것.

    출처: https://huggingface.co/ByteDance-Seed/Seed-X-PPO-7B
    """
    src = LANGUAGES[source_lang]["english"]
    tgt = LANGUAGES[target_lang]["english"]
    tag = LANGUAGES[target_lang]["seedx"]
    return f"Translate the following {src} sentence into {tgt}:\n{text} <{tag}>"


class Translator(abc.ABC):
    def __init__(self, spec: MTSpec, device: str = "cuda", device_index: int = 0) -> None:
        self.spec = spec
        self.device = device
        self.device_index = device_index
        self._lock = threading.Lock()

    @property
    @abc.abstractmethod
    def loaded(self) -> bool: ...

    @abc.abstractmethod
    def load(self) -> None: ...

    @abc.abstractmethod
    def unload(self) -> None: ...

    @abc.abstractmethod
    def _translate(
        self,
        text: str,
        source_lang: str,
        target_lang: str,
        glossary: Glossary | None,
        speech_level: SpeechLevel,
        source_politeness: Politeness,
    ) -> str: ...

    def translate(
        self,
        text: str,
        source_lang: str,
        target_lang: str,
        glossary: Glossary | None = None,
        speech_level: SpeechLevel = SpeechLevel.POLITE,
        source_politeness: Politeness = Politeness.UNKNOWN,
    ) -> str:
        """말투(speech_level)는 지시문을 이해하는 모델에만 프롬프트로 전달된다.

        나머지 백엔드는 상위에서 `apply_speech_level` 후처리로 맞춘다.
        """
        text = text.strip()
        if not text:
            return ""
        if source_lang == target_lang:
            return text
        out = self._translate(
            text, source_lang, target_lang, glossary, speech_level, source_politeness
        )
        out = _TRAILING_NOISE.sub("", out).strip()
        if target_lang == "ko":
            # 용어집이 단어를 바꿔치기하면 앞말 받침이 달라져 조사가 어긋난다.
            # (자기장를 -> 자기장을)
            out = fix_particles(out)
        return out


class CTranslate2Translator(Translator):
    """NLLB 계열. 용어집은 플레이스홀더 보호 방식으로 적용한다."""

    def __init__(self, spec: MTSpec, device: str = "cuda", device_index: int = 0) -> None:
        super().__init__(spec, device, device_index)
        self._model = None
        self._tokenizer = None

    @property
    def loaded(self) -> bool:
        return self._model is not None

    def load(self) -> None:
        if self._model is not None:
            return
        import ctranslate2
        from huggingface_hub import snapshot_download
        from transformers import AutoTokenizer

        local = snapshot_download(self.spec.repo, cache_dir=str(models_dir()))
        compute = self.spec.compute_type
        if self.device == "cpu" and "float16" in compute:
            compute = "int8"
        log.info("번역 모델 로드: %s (%s/%s)", self.spec.repo, self.device, compute)
        self._model = ctranslate2.Translator(
            local,
            device=self.device,
            device_index=self.device_index,
            compute_type=compute,
        )
        self._tokenizer = AutoTokenizer.from_pretrained(local, cache_dir=str(models_dir()))

    def unload(self) -> None:
        with self._lock:
            self._model = None
            self._tokenizer = None

    def _run_model(self, text: str, source_lang: str, target_lang: str) -> str:
        """용어집 처리를 뺀 순수 번역 한 번."""
        src_code = LANGUAGES[source_lang]["nllb"]
        tgt_code = LANGUAGES[target_lang]["nllb"]
        with self._lock:
            self._tokenizer.src_lang = src_code
            tokens = self._tokenizer.convert_ids_to_tokens(
                self._tokenizer.encode(text, truncation=True, max_length=512)
            )
            results = self._model.translate_batch(
                [tokens],
                target_prefix=[[tgt_code]],
                beam_size=2,
                max_decoding_length=512,
                repetition_penalty=1.1,
            )
            hyp = results[0].hypotheses[0]
            if hyp and hyp[0] == tgt_code:
                hyp = hyp[1:]
            return self._tokenizer.decode(
                self._tokenizer.convert_tokens_to_ids(hyp), skip_special_tokens=True
            )

    def _translate(self, text, source_lang, target_lang, glossary,
                   speech_level=SpeechLevel.POLITE, source_politeness=Politeness.UNKNOWN):
        # seq2seq 는 프롬프트가 없어 말투를 지시할 수 없다. 후처리로 맞춘다.
        if self._model is None:
            self.load()
        assert self._model is not None and self._tokenizer is not None

        if glossary is None or not len(glossary):
            return self._run_model(text, source_lang, target_lang)

        protected, replacements = glossary.protect(text, target_lang)
        if not replacements:
            return self._run_model(text, source_lang, target_lang)

        out = self._run_model(protected, source_lang, target_lang)
        return _restore_or_retranslate(
            out, replacements,
            lambda: self._run_model(text, source_lang, target_lang),
        )


class LLMTranslator(Translator):
    """Seed-X / Qwen3 등 LLM 백엔드. 용어집은 프롬프트로 직접 지시한다."""

    def __init__(self, spec: MTSpec, device: str = "cuda", device_index: int = 0,
                 lora_path: str = "") -> None:
        super().__init__(spec, device, device_index)
        self.lora_path = lora_path
        self._model = None
        self._tokenizer = None

    @property
    def loaded(self) -> bool:
        return self._model is not None

    def load(self) -> None:
        if self._model is not None:
            return
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        dtype = torch.float16 if self.device == "cuda" else torch.float32
        log.info("번역 LLM 로드: %s (%s)", self.spec.repo, self.device)
        self._tokenizer = AutoTokenizer.from_pretrained(
            self.spec.repo, cache_dir=str(models_dir())
        )
        try:
            self._model = AutoModelForCausalLM.from_pretrained(
                self.spec.repo,
                cache_dir=str(models_dir()),
                torch_dtype=dtype,
                device_map={"": self.device_index} if self.device == "cuda" else "cpu",
            )
        except (ImportError, RuntimeError, ValueError) as exc:
            if self.spec.compute_type == "int4":
                # AWQ 가중치는 전용 커널이 있어야 읽힌다. 메시지가 모호해서
                # 사용자가 원인을 못 찾는 일이 잦으므로 해결 방법을 같이 알려준다.
                raise RuntimeError(
                    f"양자화 번역 모델({self.spec.repo})을 불러오지 못했습니다.\n"
                    "AWQ 커널이 필요합니다: pip install autoawq\n"
                    "설치가 어려우면 '모델' 화면에서 다른 티어를 선택하세요.\n"
                    f"원인: {exc}"
                ) from exc
            raise
        if self.lora_path:
            from peft import PeftModel

            log.info("LoRA 어댑터 적용: %s", self.lora_path)
            self._model = PeftModel.from_pretrained(self._model, self.lora_path)
        self._model.eval()

    def unload(self) -> None:
        with self._lock:
            self._model = None
            self._tokenizer = None
            try:
                import torch

                torch.cuda.empty_cache()
            except Exception:  # noqa: BLE001
                pass

    def _build_prompt(self, text, source_lang, target_lang, glossary,
                      speech_level=SpeechLevel.POLITE,
                      source_politeness=Politeness.UNKNOWN) -> str:
        if self.spec.prompt_style is PromptStyle.SEEDX:
            # Seed-X 는 지시문을 못 알아듣는다. 말투는 후처리로만 맞춘다.
            return build_seedx_prompt(text, source_lang, target_lang)
        return self._build_instruct_prompt(
            text, source_lang, target_lang, glossary, speech_level, source_politeness
        )

    def _build_instruct_prompt(self, text, source_lang, target_lang, glossary,
                               speech_level=SpeechLevel.POLITE,
                               source_politeness=Politeness.UNKNOWN) -> str:
        """지시문을 이해하는 모델용. 가능하면 chat template 을 태운다."""
        src = LANGUAGES[source_lang]["english"]
        tgt = LANGUAGES[target_lang]["english"]
        hint = glossary.prompt_hint(text, target_lang) if glossary else ""
        parts = [
            f"Translate the following {src} text into {tgt}. "
            "It is a live spoken line from a game or broadcast, so keep the tone "
            "casual and natural. Output only the translation."
        ]
        if target_lang == "ko":
            parts.append(prompt_instruction(speech_level, source_politeness))
        if hint:
            parts.append(hint)
        parts.append(f"{src}: {text}\n{tgt}:")
        instruction = "\n\n".join(parts)

        template = getattr(self._tokenizer, "chat_template", None)
        if not template:
            return instruction
        # Qwen3 는 기본적으로 thinking 모드가 켜져 실시간 자막에는 너무 느리다.
        try:
            return self._tokenizer.apply_chat_template(
                [{"role": "user", "content": instruction}],
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False,
            )
        except TypeError:
            # enable_thinking 을 모르는 템플릿
            return self._tokenizer.apply_chat_template(
                [{"role": "user", "content": instruction}],
                tokenize=False,
                add_generation_prompt=True,
            )

    def _generate(self, prompt: str) -> str:
        """프롬프트 하나를 돌려 첫 줄만 돌려준다."""
        import torch

        with self._lock:
            inputs = self._tokenizer(prompt, return_tensors="pt").to(self._model.device)
            with torch.inference_mode():
                generated = self._model.generate(
                    **inputs,
                    max_new_tokens=256,
                    do_sample=False,
                    num_beams=1,
                    repetition_penalty=1.05,
                    pad_token_id=self._tokenizer.eos_token_id,
                )
            new_tokens = generated[0][inputs["input_ids"].shape[-1] :]
            out = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
        return out.strip().split("\n")[0]

    def _translate(self, text, source_lang, target_lang, glossary,
                   speech_level=SpeechLevel.POLITE, source_politeness=Politeness.UNKNOWN):
        if self._model is None:
            self.load()
        assert self._model is not None and self._tokenizer is not None

        def run(source: str, gloss) -> str:
            return self._generate(
                self._build_prompt(
                    source, source_lang, target_lang, gloss,
                    speech_level, source_politeness,
                )
            )

        # Seed-X 는 지시문을 못 알아들으므로 용어집을 자리표시자로 보호한다.
        # 지시문을 이해하는 모델은 프롬프트로 넣으므로 보호가 필요 없다.
        needs_placeholder = (
            self.spec.prompt_style is PromptStyle.SEEDX
            and glossary is not None
            and len(glossary)
        )
        if not needs_placeholder:
            return run(text, glossary)

        protected, replacements = glossary.protect(text, target_lang)
        if not replacements:
            return run(text, glossary)

        out = run(protected, None)
        return _restore_or_retranslate(
            out, replacements, lambda: run(text, None)
        )


def create_translator(
    spec: MTSpec, device: str = "cuda", device_index: int = 0, lora_path: str = ""
) -> Translator:
    if spec.backend is MTBackend.CTRANSLATE2:
        return CTranslate2Translator(spec, device, device_index)
    return LLMTranslator(spec, device, device_index, lora_path)