processor.py

# 文件名: processor.py

import os
import sys
import contextlib
from io import StringIO
import json
from pathlib import Path
from datetime import datetime
import re
import time
import traceback

import streamlit as st
import ebooklib
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

from book_maker.loader import BOOK_LOADER_DICT
from book_maker.translator import MODEL_DICT
from book_maker.utils import LANGUAGES, prompt_config_to_kwargs


# --- 本地模型函数 ---
@st.cache_resource
def load_helsinki_model(model_name):
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        return tokenizer, model
    except Exception as e:
        raise RuntimeError(f"Helsinki Model loading failed: {e}")

def helsinki_translate_paragraphs(paragraphs, language, direction_str, status_container):
    if not direction_str or "en-zh" in direction_str:
        model_name = "Helsinki-NLP/opus-mt-en-zh"
    elif "zh-en" in direction_str:
        model_name = "Helsinki-NLP/opus-mt-zh-en"
    else:
        raise ValueError(f"未知的 Helsinki-NLP 翻译方向: {direction_str}")

    tokenizer, model = load_helsinki_model(model_name)
    translated_paragraphs = []
    pbar_container = status_container.container()
    pbar = pbar_container.progress(0, text="Helsinki模型翻译进度: 0%")
    for i, para in enumerate(paragraphs):
        if para.strip():
            inputs = tokenizer(para, return_tensors="pt", padding=True, truncation=True, max_length=512)
            translated_tokens = model.generate(**inputs)
            translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
            translated_paragraphs.append(translated_text)
        else:
            translated_paragraphs.append("")
        progress = (i + 1) / len(paragraphs)
        if pbar:
            pbar.progress(progress, text=f"Helsinki模型翻译进度: {i+1}/{len(paragraphs)}")
    return translated_paragraphs

# --- 通用工具 ---
@contextlib.contextmanager
def st_redirect(streamlit_element):
    original_stdout = sys.stdout
    output_catcher = StringIO()
    sys.stdout = output_catcher
    try:
        yield
    finally:
        sys.stdout = original_stdout

# --- 【核心】独立的任务执行器 ---
def run_translation_task(task_id, input_file, engine_id, language, api_key, kwargs_options):
    DB_FILE = Path("tasks_db.json")
    def update_task_status(status, result_file=None, error_message=None):
        tasks = {}
        if DB_FILE.exists():
            lock_file = DB_FILE.with_suffix(".lock")
            while lock_file.exists(): time.sleep(0.1)
            try:
                lock_file.touch()
                with open(DB_FILE, "r", encoding="utf-8") as f: tasks = json.load(f)
            finally:
                if lock_file.exists(): lock_file.unlink()
        if task_id in tasks:
            tasks[task_id]["status"] = status
            tasks[task_id]["updated_at"] = datetime.now().isoformat()
            if result_file: tasks[task_id]["result_file"] = str(result_file)
            if error_message: tasks[task_id]["error"] = error_message
        lock_file = DB_FILE.with_suffix(".lock")
        while lock_file.exists(): time.sleep(0.1)
        try:
            lock_file.touch()
            with open(DB_FILE, "w", encoding="utf-8") as f: json.dump(tasks, f, indent=2, ensure_ascii=False)
        finally:
            if lock_file.exists(): lock_file.unlink()
    try:
        update_task_status("🏃‍♂️ 运行中...")
        class MockStreamlitElement:
            def empty(self): return self
            def container(self): return self
            def progress(self, *args, **kwargs): pass
            def text(self, *args, **kwargs): pass
            def info(self, *args, **kwargs): pass
        mock_status_container = MockStreamlitElement()
        output_file = translate_book_processing(input_file, engine_id, api_key, language, mock_status_container, **kwargs_options)
        update_task_status("✅ 已完成", result_file=output_file)
    except Exception as e:
        error_details = traceback.format_exc()
        update_task_status("❌ 失败", error_message=f"{str(e)}\n\nTraceback:\n{error_details}")

# --- 【核心】原始翻译处理函数 ---
def translate_book_processing(input_file_path, engine_id, api_key, language, status_container, **kwargs):
    book_type = input_file_path.split('.')[-1]
    
    # 【修复】从带有时间戳的输入文件名中，恢复出原始文件名
    # 例如：从 "100849_【初祥】发烧.txt" 变为 "【初祥】发烧.txt"
    original_name_with_ext = "_".join(Path(input_file_path).name.split('_')[1:])
    name_only, ext_only = os.path.splitext(original_name_with_ext)
    
    output_dir = Path("outputs")
    output_dir.mkdir(exist_ok=True)
    # 【修复】使用恢复后的原始文件名来创建双语文件名
    output_file_path = str(output_dir / f"{name_only}_bilingual{ext_only}")

    if engine_id == "local_helsinki":
        single_translate = kwargs.get("single_translate", False)
        book_loader_class = BOOK_LOADER_DICT.get(book_type)
        if not book_loader_class: raise ValueError(f"不支持的文件格式: {book_type}")
        temp_loader_options = {"model": lambda k, l, **kw: None, "key": "", "language": language, "resume": kwargs.get("resume", False)}
        temp_loader_options[f"{book_type}_name"] = input_file_path
        book = book_loader_class(**temp_loader_options)
        if book_type == "epub":
            new_book = book._make_new_book(book.origin_book)
            all_items = list(book.origin_book.get_items_of_type(ebooklib.ITEM_DOCUMENT))
            for item in book.origin_book.get_items():
                if item.get_type() != ebooklib.ITEM_DOCUMENT: new_book.add_item(item)
            for item in all_items:
                soup = book.bs(item.content, "html.parser")
                tags_to_translate = kwargs.get("translate_tags", "p").split(",")
                p_list = soup.findAll(tags_to_translate)
                original_texts = [p.get_text() for p in p_list]
                translated_texts = helsinki_translate_paragraphs(original_texts, language, kwargs.get("direction"), status_container)
                book.helper = book_loader_class.helper_class(translate_model=None)
                for p, trans in zip(p_list, translated_texts):
                    book.helper.insert_trans(p, trans, "", single_translate)
                item.content = soup.encode()
                new_book.add_item(item)
            ebooklib.epub.write_epub(output_file_path, new_book, {})
        else:
            with open(input_file_path, 'r', encoding='utf-8') as f: paragraphs = f.read().split('\n\n')
            translated_paragraphs = helsinki_translate_paragraphs(paragraphs, language, kwargs.get("direction"), status_container)
            with open(output_file_path, 'w', encoding='utf-8') as f:
                for orig, trans in zip(paragraphs, translated_paragraphs):
                    if not single_translate: f.write(orig + '\n\n')
                    f.write(trans + '\n\n')
        return output_file_path

    translate_model_class = MODEL_DICT.get(engine_id)
    if not translate_model_class: raise ValueError(f"不支持的翻译引擎: {engine_id}")
    loader_arg_name = f"{book_type}_name"
    language_value = LANGUAGES.get(language, language)
    loader_options = {
        loader_arg_name: input_file_path, "model": translate_model_class,
        "key": api_key or "None", "language": language_value, "resume": kwargs.get("resume", False),
        "is_test": kwargs.get("is_test", False), "test_num": kwargs.get("test_num", 10),
        "prompt_config": kwargs.get("prompt_config"), "single_translate": kwargs.get("single_translate", False),
        "context_flag": kwargs.get("context_flag", False), "temperature": kwargs.get("temperature", 1.0),
        "source_lang": kwargs.get("source_lang", "auto"), "model_api_base": kwargs.get("api_base"),
    }
    proxy = kwargs.get('proxy')
    if proxy: os.environ["http_proxy"] = os.environ["https_proxy"] = proxy
    book_translator = BOOK_LOADER_DICT.get(book_type)(**loader_options)
    with st_redirect(status_container):
        book_translator.make_bilingual_book()

    if os.path.exists(output_file_path):
        return output_file_path
    
    # Fallback check for original library's naming convention
    # e.g., for input 'uploads/123_test.txt', it might create 'uploads/123_test_bilingual.txt'
    temp_output_path = f"{os.path.splitext(input_file_path)[0]}_bilingual{ext_only}"
    if os.path.exists(temp_output_path):
        os.rename(temp_output_path, output_file_path)
        return output_file_path
        
    raise FileNotFoundError(f"翻译完成，但未在预期的输出路径找到文件: {output_file_path}")