import argparse
import sys
import os
import asyncio
from pathlib import Path
from dotenv import load_dotenv

# Add project root to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from src.common.utils import setup_logger, ensure_directory
from src.common.paths import get_work_dirs
from src.common.config import load_global_config
from src.translation.manifest_manager import ManifestManager
from src.translation.llm_client import LLMClient
from src.assembly.restoration_engine import RestorationEngine

logger = setup_logger("pipeline_restore")

async def run_restore(args):
    input_path = Path(args.input_epub)
    if not input_path.exists():
        logger.error(f"Input file not found: {input_path}")
        sys.exit(1)

    work = get_work_dirs(input_path)
    manifest_path = work["manifest"]
    
    if not manifest_path.exists():
        logger.error("Missing manifest. Run Step 1 & 2.")
        sys.exit(1)
        
    # Load Config
    load_dotenv()
    config = load_global_config()
    llm_conf = config.get("llm", {})
    api_key = llm_conf.get("api_key") or os.getenv("OPENAI_API_KEY")
    
    llm_client = None
    if api_key:
        logger.info("Initializing LLM Client for repairs...")
        llm_client = LLMClient(
            api_key=api_key,
            base_url=llm_conf.get("base_url"),
            model=llm_conf.get("model", "gpt-3.5-turbo"),
            requests_per_minute=llm_conf.get("requests_per_minute", 60),
            concurrent_requests=llm_conf.get("concurrent_requests", 5),
            chunk_dir=work["chunks"] # Reuse chunks dir for logging repairs
        )
    else:
        logger.warning("No API Key. LLM Repair disabled.")
        
    try:
        # Load Manifest
        manager = ManifestManager(manifest_path)
        manager.load()
        logger.info(f"Loaded {len(manager.entries)} entries.")
        
        # Restore Phase
        engine = RestorationEngine(llm_client)
        logger.info("Starting format restoration (Spacing + Tags + Repair)...")
        if args.force:
            logger.info("Force mode enabled: Re-processing all entries.")
            
        success_count = await engine.restore_entries(manager.entries, force=args.force)
        
        # Save Result
        manager.save()
        logger.info(f"Restoration complete. {success_count}/{len(manager.entries)} fully validated.")
        
    except Exception as e:
        logger.critical(f"Restoration failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    finally:
        if llm_client:
            await llm_client.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Step 3: Restore Format (HTML Generation)")
    parser.add_argument("input_epub", help="Path to original EPUB")
    parser.add_argument("--force", action="store_true", help="Force re-restoration")
    
    args = parser.parse_args()
    asyncio.run(run_restore(args))