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/

def get_work_dirs(input_path: Path) -> dict:
    """Get work directory paths for a specific book."""
    book_name = input_path.stem
    work_root = Path(".work") / book_name
    
    return {
        "root": work_root,
        "structure": work_root / "book_structure.json",
        "manifest": work_root / "manifest.json",
        "assets": work_root / "assets",
        "chunks": work_root / "chunks",
    }


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)
            )
            
            # 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)
            translator = Translator(llm_client, chunk_size=target_chunk_size)
            await translator.translate(manifest_manager.entries, profile)
            manifest_manager.save()
            
            await llm_client.close()
        else:
             logger.info("Skipping translation step.")
        
        # 5. Backfill
        backfiller = BackfillEngine()
        updated_structure = 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)
        output_filename = f"bilingual_{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}")
        
    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("--mode", default="bilingual", choices=["bilingual", "target_only"], help="Output mode")
    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()
    
    asyncio.run(run_pipeline(args))

if __name__ == "__main__":
    main()