import gradio as gr
import torch
import fitz  # PyMuPDF
from transformers import AutoModelForCausalLM, AutoTokenizer

# モデルID
model_id = "tencent/HY-MT1.5-1.8B"

# 環境に合わせてデバイスと精度を自動選択
if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
else:
    device = "cpu"
    dtype = torch.float32

print(f"Loading model on {device} with {dtype}...")

# トークナイザーとモデルの読み込み
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=device,
    torch_dtype=dtype
)

def extract_text_from_pdf(pdf_file):
    """PDF에서 텍스트 추출"""
    if pdf_file is None:
        return ""
    
    try:
        doc = fitz.open(pdf_file.name)
        full_text = ""
        
        for page_num, page in enumerate(doc, 1):
            text = page.get_text("text")
            if text.strip():
                full_text += f"\n--- Page {page_num} ---\n{text.strip()}\n"
        
        doc.close()
        return full_text.strip()
    
    except Exception as e:
        return f"❌ PDF 추출 오류: {str(e)}"

def translate_text(source_text, target_lang):
    """텍스트 번역"""
    if not source_text or not source_text.strip():
        return "입력 텍스트가 없습니다."
    
    # プロンプトの切り替えロジック
    if "Chinese" in target_lang or "中文" in target_lang:
        prompt = f"将以下文本翻译为{target_lang}，注意只需要输出翻译后的结果，不要额外解释：\n{source_text}"
    else:
        prompt = f"Translate the following segment into {target_lang}, without additional explanation.\n{source_text}"
    
    messages = [{"role": "user", "content": prompt}]
    
    # 入力処理
    text_input = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=False, 
        return_tensors="pt"
    ).to(device)
    
    # 生成実行
    with torch.no_grad():
        generated_ids = model.generate(
            text_input,
            max_new_tokens=1024,
            temperature=0.7,
            top_p=0.6,
            repetition_penalty=1.05
        )
    
    # 出力処理
    input_length = text_input.shape[1]
    response = generated_ids[0][input_length:]
    decoded_output = tokenizer.decode(response, skip_special_tokens=True)
    
    return decoded_output

def translate_long_text(source_text, target_lang, chunk_size=1500):
    """긴 텍스트를 청크로 나눠서 번역"""
    if not source_text or not source_text.strip():
        return "입력 텍스트가 없습니다."
    
    # 짧은 텍스트는 바로 번역
    if len(source_text) <= chunk_size:
        return translate_text(source_text, target_lang)
    
    # 긴 텍스트는 문단 단위로 분할
    paragraphs = source_text.split('\n\n')
    chunks = []
    current_chunk = ""
    
    for para in paragraphs:
        if len(current_chunk) + len(para) < chunk_size:
            current_chunk += para + "\n\n"
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = para + "\n\n"
    
    if current_chunk:
        chunks.append(current_chunk.strip())
    
    # 각 청크 번역
    translated_chunks = []
    for i, chunk in enumerate(chunks):
        print(f"Translating chunk {i+1}/{len(chunks)}...")
        translated = translate_text(chunk, target_lang)
        translated_chunks.append(translated)
    
    return "\n\n".join(translated_chunks)

def process_pdf_and_translate(pdf_file, target_lang):
    """PDF 업로드 → 텍스트 추출 → 번역"""
    if pdf_file is None:
        return "", "PDF 파일을 업로드해주세요."
    
    # 텍스트 추출
    extracted_text = extract_text_from_pdf(pdf_file)
    
    if extracted_text.startswith("❌"):
        return "", extracted_text
    
    if not extracted_text.strip():
        return "", "PDF에서 텍스트를 추출할 수 없습니다."
    
    # 번역
    translated_text = translate_long_text(extracted_text, target_lang)
    
    return extracted_text, translated_text

def translate_input_text(source_text, target_lang):
    """입력 텍스트 번역"""
    return translate_long_text(source_text, target_lang)

# UIの構築
langs = ["Japanese", "English", "Chinese", "Korean", "French", "German", "Spanish", "한국어", "日本語", "中文"]

with gr.Blocks(title="HY-MT1.5 Translator") as demo:
    gr.Markdown("# 🚀 HY-MT1.5-1.8B Translator")
    gr.Markdown("Tencent의 1.8B 번역 모델을 사용한 텍스트/PDF 번역 데모입니다.")
    
    with gr.Tabs():
        # ============ Tab 1: 텍스트 번역 ============
        with gr.TabItem("📝 Text Translation"):
            with gr.Row():
                with gr.Column():
                    input_text = gr.Textbox(
                        label="원문 (Source Text)", 
                        lines=10, 
                        placeholder="번역할 텍스트를 입력하세요..."
                    )
                    target_lang_text = gr.Dropdown(
                        choices=langs, 
                        value="English", 
                        label="번역 언어 (Target Language)"
                    )
                    translate_btn = gr.Button("🔄 번역 (Translate)", variant="primary")
                
                with gr.Column():
                    output_text = gr.Textbox(
                        label="번역 결과 (Result)", 
                        lines=10, 
                        interactive=False
                    )
            
            translate_btn.click(
                fn=translate_input_text,
                inputs=[input_text, target_lang_text],
                outputs=output_text
            )
        
        # ============ Tab 2: PDF 번역 ============
        with gr.TabItem("📄 PDF Translation"):
            gr.Markdown("### PDF 파일을 업로드하면 텍스트를 추출하고 번역합니다.")
            
            with gr.Row():
                with gr.Column():
                    pdf_input = gr.File(
                        label="📄 PDF 파일 업로드",
                        file_types=[".pdf"]
                    )
                    target_lang_pdf = gr.Dropdown(
                        choices=langs, 
                        value="English", 
                        label="번역 언어 (Target Language)"
                    )
                    translate_pdf_btn = gr.Button("🔄 PDF 번역", variant="primary")
            
            with gr.Row():
                with gr.Column():
                    extracted_text = gr.Textbox(
                        label="📋 추출된 원문 (Extracted Text)",
                        lines=15,
                        interactive=False
                    )
                
                with gr.Column():
                    translated_pdf_text = gr.Textbox(
                        label="📋 번역 결과 (Translated Text)",
                        lines=15,
                        interactive=False
                    )
            
            translate_pdf_btn.click(
                fn=process_pdf_and_translate,
                inputs=[pdf_input, target_lang_pdf],
                outputs=[extracted_text, translated_pdf_text]
            )

# 起動
demo.launch()