import gradio as gr
from huggingface_hub import InferenceClient
import os

# Cohere Command R+ 모델 ID 정의
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

def get_client():
    """
    Cohere Command R+ 모델을 위한 InferenceClient 생성.
    토큰은 환경 변수에서 가져옴.
    """
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다.")
    return InferenceClient(COHERE_MODEL, token=hf_token)


def translate_text(text, source_lang, target_lang):
    """
    텍스트를 번역하는 함수.
    """
    try:
        client = get_client()

        # 프롬프트 규칙 설정
        if source_lang == "한국어" and target_lang == "영어":
            prompt = f"""
            You are a 30-year veteran English translator.
            You are renowned for your accurate and meticulous expressions.
            Translate the following text from Korean to English:
            '{text}'
            - Translate accurately according to the context and flow of the content.
            - Do not add any extra text or explanations, just provide the translation.
            """
        elif source_lang == "영어" and target_lang == "한국어":
            prompt = f"""
            너는 30년차 한국어 전문 번역가이다.
            정확하고 섬세한 표현력으로 유명한 한국어 전문 번역가이다.
            다음 텍스트를 영어에서 한국어로 번역하라:
            '{text}'
            - 문맥과 내용의 흐름에 맞춰 정확한 번역을 한다.
            - 추가적인 설명 없이 번역 결과만 제공하라.
            """
        else:
            return "지원하지 않는 언어 조합입니다."

        # 번역에 최적화된 설정
        response = client.text_generation(
            prompt,
            max_new_tokens=200,  # 번역 결과의 길이를 적절히 제한
            temperature=0.3,     # 창의성을 낮춰 정확한 번역을 유도
            top_p=0.9            # 높은 확률의 단어만 선택하도록 설정
        )
        return response.strip()
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"


# Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("# 번역기")

    # 번역기 탭
    with gr.Tab("번역기"):
        # 언어 선택
        source_lang = gr.Radio(
            choices=["한국어", "영어"],
            label="원본 언어",
            value="한국어"
        )
        target_lang = gr.Radio(
            choices=["한국어", "영어"],
            label="목표 언어",
            value="영어"
        )

        # 입력 텍스트
        input_text = gr.Textbox(label="번역할 텍스트", lines=5)

        # 번역 결과
        translation_output = gr.Textbox(label="번역 결과", lines=5, interactive=False)

        # 번역 버튼
        translate_button = gr.Button("번역")

        # 번역 함수 연결
        translate_button.click(
            fn=translate_text,
            inputs=[input_text, source_lang, target_lang],
            outputs=translation_output
        )

# 메인 실행부
if __name__ == "__main__":
    demo.launch()