from logging import getLogger
from pathlib import Path
from llama_cpp import Llama
from functools import lru_cache

from lib.utils import Timer

logger = getLogger(__name__)
LLM_SYS_PROMPT_EN= """
你是一名专业的同声传译员，正在为 GOSIM 会议提供英中翻译服务。你的任务是准确、流畅地翻译发言内容。

请遵循以下要求：
1. 语言风格：翻译成中文时，请使用自然、流畅、符合现代汉语口语习惯的表达方式。避免生硬、逐字翻译的痕迹，要让听众容易理解。
2. 专业术语：**请优先参考下方提供的术语对照表进行翻译。** 对于对照表中未包含的术语，如果该术语有公认的标准翻译，请使用标准翻译；如果没有或不确定，请保留英文原文。不要用通俗词汇替代专业术语。
3. 专有名词：对于专有名词，如会议名称 "GOSIM"、人名、公司名、项目名、特定技术名称等，请保留其原始英文不做翻译。
4. 流畅性与准确性：在追求口语化的同时，务必保证信息传达的准确性。
5. 输出：请直接输出翻译结果，不要添加任何额外的解释或说明。

**专业术语对照表：**
* driver: 驱动
* bus: 总线
* mask: 掩码
* preemption: 抢占
* register: 寄存器
* Library: 库
* biases: 偏移
* OpenAGI: OpenAGI
* LLaMA Factory: LLaMA Factory
* OPENGL: OPENGL

现在，请将以下内容翻译成中文：
"""

LLM_SYS_PROMPT_ZH = """
你是一位中英文翻译专家。请将以下中文文本翻译成英文，遵循以下要求：

翻译要求：
- 保留原文英文内容：以下内容请保持原始英文形式，不进行翻译或改写：
- 技术术语与专业词汇
- 产品名称、品牌名称
- 代码片段、函数名、变量名
- 专有名词、缩写、首字母缩略词（如 API、NLP、RAG 等）
- 翻译符合英文表达习惯，流畅自然，不生硬直译。
- 保持专业性与准确性，清晰传达原意。
- 如遇原文表达模糊或逻辑不清的情况，允许适度调整语序或措辞，以增强英文表述的清晰度和逻辑性。

注意：
若难以确定某个词汇是否需要翻译，请优先保留原始英文形式。
不需添加额外解释或注释，仅翻译正文内容。
特别注意，翻译的内容只能包含英文，不能包含其他的语言。

文本："""
MODEL_PATH = Path("/Users/jeqin/work/code/Translator/python_server/moyoyo_asr_models/qwen2.5-1.5b-instruct-q5_0.gguf")
class QwenTranslator:
    def __init__(self, model_path=MODEL_PATH, system_prompt_en=LLM_SYS_PROMPT_EN, system_prompt_zh=LLM_SYS_PROMPT_ZH) -> None:
        with Timer("load llm"):
            self.llm = Llama(
            model_path=str(model_path),
            chat_format="chatml",
            n_gpu_layers=-1,
            verbose=False)
        self.sys_prompt_en = system_prompt_en
        self.sys_prompt_zh = system_prompt_zh
        self._warmup()

    def to_message(self, prompt, src_lang, dst_lang):
        """构造提示词"""
        return [
            {"role": "system", "content": self.sys_prompt_en if src_lang == "en" else self.sys_prompt_zh},
            {"role": "user", "content": prompt},
        ]

    def _warmup(self):
        self.translate(prompt="how are you?", src_lang="en", dst_lang="zh")

    @lru_cache(maxsize=10)
    def translate(self, prompt, src_lang, dst_lang):
        message = self.to_message(prompt, src_lang, dst_lang)
        with Timer("llm inference") as t:
            output = self.llm.create_chat_completion(messages=message, temperature=0)
        return output['choices'][0]['message']['content'], t.duration


if __name__ == '__main__':
    model_dir = Path("/Users/jeqin/work/code/Translator/moyoyo_asr_models")
    qwen2 = (model_dir / "qwen2.5-1.5b-instruct-q5_0.gguf").as_posix()
    qwen3 = (model_dir / "Qwen_Qwen3-1.7B-Q4_K_M.gguf").as_posix()

    translator = QwenTranslator(qwen3)
    text, time_cost =translator.translate("今天天气怎么样？", "zh", "en")
    print(text)
    print(time_cost)