#!/usr/bin/env python3
"""
Translation Pipeline Debug Script
Shows visible results at each step of the translation process.
"""

import asyncio
import json
import os
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv

from src.manifest_manager import ManifestManager
from src.llm_client import LLMClient
from src.book_profiler import BookProfiler
from src.translator import Translator
from src.format_restorer import FormatRestorer
from src.data_model import ManifestEntry, BookProfile

load_dotenv()

# Configuration
MANIFEST_PATH = Path("cache/Empire of AI Dreams and Nightmares in Sam Altmans OpenAI (Karen Hao)_manifest.json")
API_KEY = os.getenv("V3_API_KEY")
BASE_URL = "https://api.gpt.ge/v1"
MODEL = "gemini-3-flash-preview"
EXTRA_HEADERS = {"x-foo": "true"}

# Chunk size configuration
MAX_CHUNK_SIZE = 15  # Maximum entries per chunk


def group_entries_by_file(entries: list) -> dict:
    """Group manifest entries by their source file (chapter)."""
    grouped = {}
    for entry in entries:
        file_path = entry.file_path
        if file_path not in grouped:
            grouped[file_path] = []
        grouped[file_path].append(entry)
    return grouped


def create_chapter_aware_chunks(entries: list, max_size: int = MAX_CHUNK_SIZE) -> list:
    """
    Create chunks that respect chapter boundaries.
    Returns list of (file_path, chunk_entries) tuples.
    """
    grouped = group_entries_by_file(entries)
    chunks = []
    
    for file_path, file_entries in grouped.items():
        # Split this file's entries into chunks of max_size
        for i in range(0, len(file_entries), max_size):
            chunk = file_entries[i:i + max_size]
            chunks.append((file_path, chunk))
    
    return chunks


def get_file_type(file_path: str) -> str:
    """Determine the type of content based on file name."""
    fname = file_path.lower()
    if any(k in fname for k in ['toc', 'contents', 'nav']):
        return 'toc'
    elif any(k in fname for k in ['title', 'cover']):
        return 'cover'
    elif any(k in fname for k in ['copyright', 'colophon']):
        return 'legal'
    elif any(k in fname for k in ['author', 'about']):
        return 'author_bio'
    elif any(k in fname for k in ['index', 'bibliography', 'endnote', 'footnote']):
        return 'reference'
    else:
        return 'body'


async def step1_load_manifest():
    """Step 1: Load manifest and show statistics."""
    print("\n" + "="*60)
    print("STEP 1: Loading Manifest")
    print("="*60)
    
    manager = ManifestManager(MANIFEST_PATH)
    manager.load()
    
    entries = manager.entries
    untranslated = [e for e in entries if not e.translated_text]
    
    print(f"  Total entries: {len(entries)}")
    print(f"  Untranslated:  {len(untranslated)}")
    
    # Show grouping by file
    grouped = group_entries_by_file(entries)
    print(f"  Unique files:  {len(grouped)}")
    
    # Show sample entry
    if entries:
        sample = entries[0]
        print(f"\n  Sample entry:")
        print(f"    ID:       {sample.entry_id}")
        print(f"    File:     {sample.file_path}")
        print(f"    Original: {sample.original_text[:80]}...")
        print(f"    Placeholders: {sample.placeholders}")
    
    return manager


async def step2_profile_book(manager: ManifestManager, llm_client: LLMClient):
    """Step 2: Generate book profile."""
    print("\n" + "="*60)
    print("STEP 2: Generating Book Profile")
    print("="*60)
    
    profiler = BookProfiler(llm_client)
    profile = await profiler.analyze(manager.entries)
    
    print(f"  Title:       {profile.title}")
    print(f"  Author:      {profile.author}")
    print(f"  Genre:       {profile.genre}")
    print(f"  Keywords:    {profile.keywords}")
    print(f"  Style Guide: {profile.style_guide[:200]}..." if profile.style_guide else "  Style Guide: (none)")
    
    return profile


async def step3_create_chunks(manager: ManifestManager):
    """Step 3: Create chapter-aware chunks."""
    print("\n" + "="*60)
    print("STEP 3: Creating Chapter-Aware Chunks")
    print("="*60)
    
    untranslated = [e for e in manager.entries if not e.translated_text]
    chunks = create_chapter_aware_chunks(untranslated)
    
    print(f"  Total chunks: {len(chunks)}")
    
    # Show chunk distribution
    print(f"\n  Chunk distribution by file type:")
    type_counts = {}
    for file_path, chunk_entries in chunks:
        ftype = get_file_type(file_path)
        type_counts[ftype] = type_counts.get(ftype, 0) + 1
    
    for ftype, count in sorted(type_counts.items()):
        print(f"    {ftype}: {count} chunks")
    
    # Show first 3 chunks
    print(f"\n  First 3 chunks:")
    for i, (file_path, chunk_entries) in enumerate(chunks[:3]):
        ftype = get_file_type(file_path)
        print(f"    [{i}] {file_path} ({ftype}): {len(chunk_entries)} entries")
        if chunk_entries:
            print(f"        First: {chunk_entries[0].original_text[:50]}...")
    
    return chunks


async def step4_translate_sample(chunks: list, llm_client: LLMClient, profile: BookProfile):
    """Step 4: Translate a sample chunk and show results."""
    print("\n" + "="*60)
    print("STEP 4: Translating Sample Chunk")
    print("="*60)
    
    if not chunks:
        print("  No chunks to translate!")
        return
    
    # Pick a proper body chapter (skip first few files which are usually cover/copyright/toc)
    sample_chunk = None
    skip_prefixes = ['cM', 'c9', 'c18']  # Cover, title, contents pages
    for file_path, chunk_entries in chunks:
        # Skip non-body files and known cover/toc files
        ftype = get_file_type(file_path)
        fname = Path(file_path).stem
        if ftype == 'body' and fname not in skip_prefixes and len(chunk_entries) > 3:
            sample_chunk = (file_path, chunk_entries[:5])  # Limit to 5 entries for demo
            break
    
    if not sample_chunk:
        # Fallback to any body chunk
        for file_path, chunk_entries in chunks:
            if get_file_type(file_path) == 'body':
                sample_chunk = (file_path, chunk_entries[:5])
                break
    
    if not sample_chunk:
        sample_chunk = chunks[0]
        sample_chunk = (sample_chunk[0], sample_chunk[1][:5])
    
    file_path, entries = sample_chunk
    ftype = get_file_type(file_path)
    
    print(f"  Selected chunk: {file_path} ({ftype})")
    print(f"  Entries: {len(entries)}")
    
    # Show entries before translation
    print(f"\n  === Before Translation ===")
    for i, entry in enumerate(entries):
        print(f"  [{i}] {entry.entry_id}")
        print(f"      Original: {entry.original_text[:60]}...")
        if entry.placeholders:
            print(f"      Placeholders: {list(entry.placeholders.keys())}")
    
    # Translate
    print(f"\n  Translating...")
    results = await llm_client.translate_chunk(
        entries,
        instruction=profile.style_guide,
        mode="bilingual"
    )
    
    # Apply results and show
    print(f"\n  === After Translation ===")
    restorer = FormatRestorer()
    for i, entry in enumerate(entries):
        if entry.entry_id in results:
            translated = results[entry.entry_id]
            entry.translated_text = translated
            
            print(f"  [{i}] {entry.entry_id}")
            print(f"      Original:    {entry.original_text[:50]}...")
            print(f"      Translated:  {translated[:50]}...")
            
            # Restore format
            if entry.placeholders:
                restored, success = restorer.restore(translated, entry.placeholders)
                print(f"      Restored OK: {success}")
                if not success:
                    print(f"      Restored:    {restored[:50]}...")
        else:
            print(f"  [{i}] MISSING: {entry.entry_id}")
    
    return entries


async def step5_test_placeholders(manager: ManifestManager, llm_client: LLMClient, profile: BookProfile):
    """Step 5: Test placeholder handling with entries that have placeholders."""
    print("\n" + "="*60)
    print("STEP 5: Testing Placeholder Handling")
    print("="*60)
    
    # Find entries with placeholders
    entries_with_ph = [e for e in manager.entries if e.placeholders and len(e.placeholders) > 1]
    
    print(f"  Entries with placeholders: {len(entries_with_ph)}")
    
    if not entries_with_ph:
        print("  No entries with placeholders found!")
        return
    
    # Pick 5 diverse samples
    samples = entries_with_ph[:5]
    
    print(f"\n  === Selected Samples ({len(samples)}) ===")
    for i, entry in enumerate(samples):
        ph_keys = [k for k in entry.placeholders.keys() if not k.startswith('_')]
        print(f"  [{i}] {entry.entry_id}")
        print(f"      Original: {entry.original_text[:60]}...")
        print(f"      Placeholders: {ph_keys}")
    
    # Translate
    print(f"\n  Translating {len(samples)} entries with placeholders...")
    results = await llm_client.translate_chunk(
        samples,
        instruction=profile.style_guide,
        mode="bilingual"
    )
    
    # Show results with restoration
    print(f"\n  === Translation Results ===")
    restorer = FormatRestorer()
    success_count = 0
    
    for i, entry in enumerate(samples):
        print(f"\n  [{i}] {entry.entry_id}")
        print(f"      Original:     {entry.original_text[:50]}...")
        
        if entry.entry_id in results:
            translated = results[entry.entry_id]
            print(f"      Translated:   {translated[:50]}...")
            
            # Check if placeholders are preserved
            ph_keys = [k for k in entry.placeholders.keys() if not k.startswith('_')]
            preserved = all(f"φ{k}φ" in translated or f"φ/{k}φ" in translated for k in ph_keys if k.isdigit())
            print(f"      PH Preserved: {preserved}")
            
            # Restore format
            restored, success = restorer.restore(translated, entry.placeholders)
            print(f"      Restore OK:   {success}")
            
            if success:
                success_count += 1
            else:
                print(f"      Restored:     {restored[:50]}...")
        else:
            print(f"      MISSING from results!")
    
    print(f"\n  Summary: {success_count}/{len(samples)} restored successfully")


async def main():
    """Run all steps."""
    print("="*60)
    print("TRANSLATION PIPELINE DEBUG")
    print("="*60)
    
    if not API_KEY:
        print("ERROR: V3_API_KEY not found in .env")
        return
    
    if not MANIFEST_PATH.exists():
        print(f"ERROR: Manifest not found at {MANIFEST_PATH}")
        print("Run the main pipeline first to generate the manifest.")
        return
    
    # Initialize LLM client
    llm_client = LLMClient(
        api_key=API_KEY,
        base_url=BASE_URL,
        model=MODEL,
        extra_headers=EXTRA_HEADERS
    )
    
    try:
        # Step 1: Load manifest
        manager = await step1_load_manifest()
        
        # Step 2: Profile book
        profile = await step2_profile_book(manager, llm_client)
        
        # Step 3: Create chunks
        chunks = await step3_create_chunks(manager)
        
        # Step 4: Translate sample (simple text)
        await step4_translate_sample(chunks, llm_client, profile)
        
        # Step 5: Test placeholders
        await step5_test_placeholders(manager, llm_client, profile)
        
        print("\n" + "="*60)
        print("DEBUG COMPLETE")
        print("="*60)
        
    finally:
        await llm_client.close()


if __name__ == "__main__":
    asyncio.run(main())