"""
EPUB Translator Core Module - v0.05 (Manifest & Arena)
"""

import asyncio
import os
import sys
import json
import traceback
from typing import List, Dict, Any
from pathlib import Path
from loguru import logger
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn

from .epub_parser import EPUBParser
from .llm_client import OpenRouterClient
from .text_processor import TextProcessor
from .bilingual_builder import BilingualEPUBBuilder
from .manifest_manager import ManifestManager
from .book_profiler import BookProfiler


class EPUBTranslator:
    
    def __init__(self, config: Dict, use_cache: bool = True):
        self.config = config
        self.console = Console()
        self.use_cache = use_cache
        
        self.parser = None
        self.llm_client = OpenRouterClient(config)
        self.text_processor = TextProcessor(config)
        self.profiler = BookProfiler(config, self.llm_client)
        
        self.manifest_dir = Path("cache/manifests")
        self.manifest_dir.mkdir(parents=True, exist_ok=True)

    async def translate_epub(self, epub_path: str, test_mode: bool = False, output_dir: str = None) -> str:
        try:
            epub_path = Path(epub_path)
            self.parser = EPUBParser(str(epub_path))
            
            # 1. Prepare Manifest
            manifest_path = self.manifest_dir / f"{epub_path.stem}_manifest.json"
            manifest = ManifestManager(str(manifest_path))
            
            if not manifest.load() or not self.use_cache:
                self.console.print("[yellow]Initializing Manifest...[/yellow]")
                manifest.init_manifest(book_id=epub_path.name, metadata=self.parser.get_book_info())
                content_items = self.parser.extract_all_content_items()
                for item in content_items:
                    self.text_processor.extract_to_manifest(item['content'], item['file_name'], manifest)
                manifest.save()
            
            stats = manifest.stats
            self.console.print(f"[green]Manifest loaded: {stats['total']} paragraphs[/green]")

            # 2. Profile & Arena
            profile = {}
            best_model = self.config['llm']['models']['fast']
            
            if not test_mode:
                # A. Profile
                self.console.print("[yellow]Generating Book Profile...[/yellow]")
                profile = await self.profiler.analyze_book(manifest)
                self.console.print(f"Genre: {profile.get('genre')} | Style: {profile.get('style')}")
                
                # B. Arena
                self.console.print("[yellow]Running Model Arena...[/yellow]")
                best_model = await self.profiler.run_arena(manifest, profile)
                self.console.print(f"[bold green]🏆 Winner: {best_model}[/bold green]")

            # 3. Translate
            if test_mode:
                pending = manifest.get_items(status="pending")[:5]
                if pending:
                    results = await self.llm_client.translate_chunk(
                        pending, 
                        glossary=profile.get('glossary'), 
                        instruction=profile.get('translation_instruction'),
                        model_id_override=best_model
                    )
                    for pid, trans in results.items():
                        self.console.print(f"\n[cyan]{pid}[/cyan]: {trans}")
                return "test_mode_done"

            chunks = self.text_processor.create_chunks_from_manifest(manifest)
            if chunks:
                await self._translate_concurrently(chunks, manifest, profile, best_model)
            
            # 4. Build
            self.console.print("\n[yellow]Building bilingual EPUB...[/yellow]")
            output_path = output_dir or self.config['output']['output_dir']
            builder = BilingualEPUBBuilder(self.parser.book, self.config)
            
            translation_map = {item.global_id: item.translation for item in manifest.get_items() if item.translation}
            paragraph_map = {item.global_id: {
                "file_name": item.source_file,
                "text": item.clean_text,
                "html_element": item.original_html
            } for item in manifest.get_items()}
            
            result_file = builder.create_bilingual_epub_with_mapping(
                translation_map, paragraph_map, output_path
            )
            
            self.console.print(f"[green]✅ Translation complete! File: {result_file}[/green]")
            return result_file
        except Exception as e:
            traceback.print_exc()
            logger.error(f"Translation flow failed: {e}")
            raise

    async def _translate_concurrently(self, chunks: List[List[Any]], manifest: ManifestManager, profile: Dict, model_id: str):
        total_chunks = len(chunks)
        glossary = profile.get('glossary', {})
        instruction = profile.get('translation_instruction', "")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeElapsedColumn(),
            console=self.console
        ) as progress:
            task_id = progress.add_task(f"[cyan]Translating ({model_id.split('/')[-1]})...", total=total_chunks)
            semaphore = self.llm_client.rate_limiter.semaphore
            
            async def worker(chunk, idx):
                async with semaphore:
                    try:
                        results = await self.llm_client.translate_chunk(
                            chunk, 
                            glossary=glossary, 
                            instruction=instruction,
                            model_id_override=model_id
                        )
                        for item in chunk:
                            if item.global_id in results:
                                manifest.update_item(item.global_id, results[item.global_id], model=model_id)
                            else:
                                manifest.update_item(item.global_id, None, status="failed", error="Missing")
                        manifest.save()
                    except Exception as e:
                        logger.error(f"Chunk {idx} failed: {e}")
                    finally:
                        progress.update(task_id, advance=1)

            tasks = [worker(chunk, i) for i, chunk in enumerate(chunks)]
            await asyncio.gather(*tasks)