epub_bilingual_translator/main.py

import argparse
import sys
import os
import asyncio
from pathlib import Path
from dotenv import load_dotenv
from src.common.config import load_global_config

from src.preprocessing.epub_cleaner import EpubCleaner
from src.preprocessing.profiler import BookProfiler
from src.preprocessing.text_extractor import FineGrainedExtractor
from src.translation.manifest_manager import ManifestManager
from src.translation.translator_engine import Translator
from src.translation.llm_client import LLMClient
from src.assembly.backfiller import BackfillEngine
from src.assembly.builder import BilingualBuilder
from src.common.data_model import BookStructure
from src.common.utils import setup_logger, ensure_directory
from src.common.exceptions import EpubTranslatorError

logger = setup_logger("main")

# Unified work directory structure
# .work/
#   ├── {book_name}/
#   │   ├── book_structure.json
#   │   ├── manifest.json
#   │   ├── assets/
#   │   └── chunks/

from src.common.paths import get_work_dirs


async def run_pipeline(args):
    input_path = Path(args.input_epub)
    output_dir = Path(args.output_dir)
    
    # Get work directories for this book
    work = get_work_dirs(input_path)
    ensure_directory(work["root"])
    ensure_directory(output_dir)
    
    # Load Config
    config = load_global_config()
    llm_conf = config.get("llm", {})
    trans_conf = config.get("translation", {})
    
    api_key = llm_conf.get("api_key")
    # Base URL and Model come from config if not overridden
    base_url = llm_conf.get("base_url")
    # CLI model arg overrides config model, which overrides default
    model = args.model if args.model else llm_conf.get("model", "gpt-3.5-turbo")
    
    if not api_key:
        logger.warning("OPENAI_API_KEY not found in env or config. LLM features may fail.")

    try:
        # 1. Preprocessing - Reuse book_structure.json if exists
        if work["structure"].exists() and not args.force_clean:
            logger.info(f"Reusing existing book_structure: {work['structure']}")
            structure = BookStructure.load(work["structure"])
        else:
            logger.info("Cleaning EPUB and generating book_structure...")
            cleaner = EpubCleaner(input_path, work["root"])
            book_structure_json = cleaner.clean()
            structure = BookStructure.load(book_structure_json)

        # 2. Extraction
        extractor = FineGrainedExtractor()
        manifest_entries = extractor.extract(structure)
        
        # 3. Manifest Management
        manifest_manager = ManifestManager(work["manifest"])
        manifest_manager.load()  # Load existing if any
        manifest_manager.add_entries(manifest_entries)
        manifest_manager.save()
        
        # 4. Translation
        if not args.skip_translation:
            if not api_key:
                logger.error("Cannot translate without API Key. Use --skip-translation to test pipeline.")
                sys.exit(1)
                
            llm_client = LLMClient(
                api_key=api_key, 
                base_url=base_url, 
                model=model,
                requests_per_minute=llm_conf.get("requests_per_minute", 60),
                concurrent_requests=llm_conf.get("concurrent_requests", 5),
                chunk_dir=work["chunks"]
            )
            
            # Profiling
            profiler = BookProfiler(llm_client)
            profile = await profiler.analyze(manifest_manager.entries)
            logger.info(f"Book Profile: {profile}")
            
            # Translation
            target_chunk_size = trans_conf.get("chunk_size", 5000)
            concurrent_reqs = llm_conf.get("concurrent_requests", 5)
            translator = Translator(llm_client, chunk_size=target_chunk_size, max_concurrent=concurrent_reqs)
            await translator.translate(manifest_manager.entries, profile)
            manifest_manager.save()
            
            # Don't close here, wait until after backfill
            # await llm_client.close() 
            pass
        else:
             logger.info("Skipping translation step.")
        
        # 5. Backfill (now async + LLM repair enabled)
        # Reuse existing llm_client if available, otherwise create temporary one if needed?
        # In this flow, llm_client is created inside the 'if not args.skip_translation' block.
        # If skip_translation is True, llm_client is undefined. 

        backfill_llm_client = None
        should_close_client = False
        
        if 'llm_client' in locals() and llm_client:
             backfill_llm_client = llm_client
        elif api_key and not args.skip_translation:
             # This case shouldn't happen because if not skip, we key llm_client above.
             # But if skip_translation is True, we might still want repair?
             # For now, let's only enable repair if translation occurred or if we explicitly create one.
             # User said: "LLM features may fail" if no key.
             pass
             
        # Initialization
        backfiller = BackfillEngine(llm_client=backfill_llm_client)
        updated_structure = await backfiller.backfill(structure, manifest_manager.entries, mode=args.mode)
        
        # 6. Assembly - Pass original EPUB for TOC preservation
        builder = BilingualBuilder(work["root"], original_epub_path=input_path)
        
        if args.mode == "bilingual":
            output_filename = f"bilingual_{input_path.name}"
        else:
            output_filename = f"translated_{input_path.name}"
            
        output_path = output_dir / output_filename
        
        created_epub = builder.build(updated_structure, output_path)
        
        logger.info(f"Pipeline completed! Output: {created_epub}")
        
        if 'llm_client' in locals() and llm_client:
            await llm_client.close()
        
    except EpubTranslatorError as e:
        logger.error(f"An error occurred: {e}")
        sys.exit(1)
    except Exception as e:
        logger.critical(f"Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description="EPUB Bilingual Translator")
    parser.add_argument("input_epub", help="Path to the input EPUB file")
    parser.add_argument("--output-dir", default="output", help="Directory for output files")
    parser.add_argument("--model", default=None, help="LLM Model to use (overrides config)")
    parser.add_argument("--bilingual", action="store_true", help="Output bilingual version (default is target language only)")
    parser.add_argument("--skip-translation", action="store_true", help="Skip LLM translation (for testing)")
    parser.add_argument("--force-clean", action="store_true", help="Force re-clean EPUB even if book_structure exists")
    
    args = parser.parse_args()
    
    # Map boolean flag to mode string
    args.mode = "bilingual" if args.bilingual else "target_only"
    
    asyncio.run(run_pipeline(args))

if __name__ == "__main__":
    main()