#!/usr/bin/env python3
"""
Translation Task: Multi-language pairs
This module provides a class for generating prompt variations for translation tasks.
"""

from typing import Dict, Any, List
import argparse
import pandas as pd
from pathlib import Path
import os
import sys
from datasets import load_dataset

from promptsuite.core import PROMPT_FORMAT_VARIATIONS, FEW_SHOT_KEY
from promptsuite.core.template_keys import (
    INSTRUCTION, PROMPT_FORMAT, QUESTION_KEY, GOLD_KEY,
    PARAPHRASE_WITH_LLM, FORMAT_STRUCTURE_VARIATION, TYPOS_AND_NOISE_VARIATION, INSTRUCTION_VARIATIONS
)
from base_task import BaseTask

from promptsuite_tasks.constants import (
    DEFAULT_VARIATIONS_PER_FIELD, DEFAULT_PLATFORM, DEFAULT_MODEL_NAME,
    DEFAULT_MAX_VARIATIONS_PER_ROW, DEFAULT_MAX_ROWS, DEFAULT_RANDOM_SEED
)


class TranslationTask(BaseTask):
    """Task for generating translation prompt variations."""
    
    def __init__(self, language_pair: str, variations_per_field: int = DEFAULT_VARIATIONS_PER_FIELD, api_platform: str = DEFAULT_PLATFORM, model_name: str = DEFAULT_MODEL_NAME,
                 max_rows: int = DEFAULT_MAX_ROWS, max_variations_per_row: int = DEFAULT_MAX_VARIATIONS_PER_ROW, random_seed: int = DEFAULT_RANDOM_SEED):
        self.language_pair = language_pair
        # Convert language pair name for display (e.g., cs-en -> Czech to English)
        display_pair = TranslationTask._get_display_name(language_pair)
        task_name = f"Translation Task: {display_pair}"
        output_filename = f"translation_{language_pair}_variations.json"
        
        super().__init__(
            task_name=task_name,
            output_filename=output_filename,
            variations_per_field=variations_per_field,
            api_platform=api_platform,
            model_name=model_name,
            max_rows=max_rows,
            max_variations_per_row=max_variations_per_row,
            random_seed=random_seed
        )
    
    @staticmethod
    def _get_display_name(language_pair: str) -> str:
        """Convert language pair code to display name."""
        language_names = {
            'cs': 'Czech',
            'de': 'German', 
            'fr': 'French',
            'hi': 'Hindi',
            'ru': 'Russian',
            'en': 'English'
        }
        
        if '-' in language_pair:
            source, target = language_pair.split('-')
            source_name = language_names.get(source, source.upper())
            target_name = language_names.get(target, target.upper())
            return f"{source_name} to {target_name}"
        return language_pair

    def load_data(self) -> None:
        """Load translation data using datasets library - both train and test splits."""
        display_pair = TranslationTask._get_display_name(self.language_pair)
        print(f"Loading WMT14 data for {display_pair} ({self.language_pair})")
        
        # Determine the WMT14 config to use
        # WMT14 only has X->EN configs, so for EN->X we load the reverse and swap
        if self.language_pair.startswith('en-'):
            # For EN->X, load X->EN and we'll reverse it in post_process
            reverse_pair = self.language_pair.split('-')[1] + '-en'  # en-de -> de-en
            wmt14_config = reverse_pair
            self.is_reversed = True
            print(f"Loading reverse config {wmt14_config} for {display_pair}")
        else:
            # For X->EN, load directly
            wmt14_config = self.language_pair
            self.is_reversed = False
        
        try:
            # Load both train and test splits, similar to qa_task.py
            print("Loading WMT14 train dataset...")
            train_ds = load_dataset("wmt14", wmt14_config, split="train[:50]")
            train_df = pd.DataFrame(train_ds)
            train_df['split'] = 'train'
            print(f"✅ Loaded {len(train_df)} train rows")

            print("Loading WMT14 test dataset...")
            test_ds = load_dataset("wmt14", wmt14_config, split="test[:50]")
            test_df = pd.DataFrame(test_ds)
            test_df['split'] = 'test'
            print(f"✅ Loaded {len(test_df)} test rows")

            # Combine
            df = pd.concat([train_df, test_df], ignore_index=True)
            print(f"✅ Combined total: {len(df)} rows")

            # Load into promptsuite
            self.ps.load_dataframe(df)
            print(f"✅ Loaded WMT14 dataset for {display_pair}: {len(self.ps.data)} rows")
            
        except Exception as e:
            print(f"❌ Error loading WMT14 dataset for {wmt14_config}: {e}")
            raise ValueError(f"Failed to load WMT14 data for language pair: {self.language_pair}")
        
        # Post-process the data to create language-specific columns
        self.post_process()
        
        train_count = sum(1 for split in self.ps.data['split'] if split == 'train')
        test_count = sum(1 for split in self.ps.data['split'] if split == 'test')
        print(f"✅ Data processed for {display_pair} ({train_count} train, {test_count} test)")

    def post_process(self) -> None:
        """Create language-specific columns from WMT14 translation data."""
        # Extract source and target language codes from language_pair (e.g., "de-en" -> "de", "en")
        source_code, target_code = self.language_pair.split('-')
        
        # WMT14 data structure: DataFrame with 'translation' column containing dicts
        # Each dict has language codes as keys (e.g., {'de': 'text', 'en': 'text'})
        if hasattr(self, 'is_reversed') and self.is_reversed:
            # For EN->X pairs, we loaded X->EN and need to map correctly
            # We loaded hi-en but want en-hi, so the mapping should be straightforward:
            # source_code (en) should get English text, target_code (hi) should get Hindi text
            self.ps.data[source_code] = [row[source_code] for row in self.ps.data['translation']]  # en gets en text
            self.ps.data[target_code] = [row[target_code] for row in self.ps.data['translation']]  # hi gets hi text
        else:
            # For X->EN pairs, use directly
            self.ps.data[source_code] = [row[source_code] for row in self.ps.data['translation']]
            self.ps.data[target_code] = [row[target_code] for row in self.ps.data['translation']]
        
        # Split column is already created during load_data() - no need to create it artificially

    def get_template(self) -> Dict[str, Any]:
        """Get template configuration for translation task."""
        # Dynamic template based on language pair
        display_pair = TranslationTask._get_display_name(self.language_pair)
        # Extract source and target language names and codes
        source_lang, target_lang = display_pair.split(' to ')
        source_code, target_code = self.language_pair.split('-')
        return {
            INSTRUCTION: f"You are a professional translator. Translate the following text from {display_pair}.",
            INSTRUCTION_VARIATIONS: [PARAPHRASE_WITH_LLM],
            PROMPT_FORMAT: f"{source_lang}: {{{source_code}}}\n{target_lang}: {{{target_code}}}",
            PROMPT_FORMAT_VARIATIONS: [FORMAT_STRUCTURE_VARIATION, TYPOS_AND_NOISE_VARIATION],
            GOLD_KEY: target_code,  # Target language code is the gold standard
            FEW_SHOT_KEY: {
                'count': 5,  # Number of few-shot examples
                'format': 'different_examples__different_order_per_variation',
                'split': 'train'  # Use training split for few-shot examples
            },
        }


def get_available_language_pairs() -> List[str]:
    """Get list of available language pairs from WMT14 dataset."""
    # Selected language pairs: Czech, Russian, Hindi (both directions)
    selected_pairs = [
        # To English
        "cs-en",  # Czech-English
        "hi-en",  # Hindi-English
        "ru-en",  # Russian-English
        # From English
        "en-cs",  # English-Czech
        "en-hi",  # English-Hindi
        "en-ru",  # English-Russian
    ]
    
    return sorted(selected_pairs)


def generate_all_language_pairs(variations_per_field, api_platform, model_name, max_rows, max_variations_per_row, random_seed):
    """Generate variations for all language pairs separately."""
    # Create output directory
    output_dir = Path(__file__).parent.parent / "data" / "translation"
    output_dir.mkdir(parents=True, exist_ok=True)
    
    language_pairs = get_available_language_pairs()
    print(f"🎯 Found {len(language_pairs)} language pairs:")
    for i, pair in enumerate(language_pairs, 1):
        display_name = TranslationTask._get_display_name(pair)
        print(f"  {i:2d}. {display_name} ({pair})")
    
    generated_files = []
    
    for language_pair in language_pairs:
        display_name = TranslationTask._get_display_name(language_pair)
        print(f"\n🌍 Processing language pair: {display_name}")
        print("=" * 50)
        
        try:
            task = TranslationTask(
                language_pair=language_pair,
                variations_per_field=variations_per_field,
                api_platform=api_platform,
                model_name=model_name,
                max_rows=max_rows,
                max_variations_per_row=max_variations_per_row,
                random_seed=random_seed
            )
            
            # Override output path to save in translation folder
            output_file = output_dir / f"translation_{language_pair}_variations.json"
            
            # Generate using the base class method but with custom output path
            print(f"🚀 Starting {task.task_name}")
            print("=" * 60)
            print("\n1. Loading data...")
            task.load_data()
            print("\n2. Setting up template...")
            template = task.get_template()
            task.ps.set_template(template)
            print("✅ Template configured")
            print(f"\n3. Configuring generation...")
            print(f"   Variations per field: {task.variations_per_field}")
            print(f"   API Platform: {task.api_platform}")
            print(f"   Model: {task.model_name}")
            print(f"   Max rows: {task.max_rows}")
            print(f"   Max variations per row: {task.max_variations_per_row}")
            print(f"   Random seed: {task.random_seed}")
            task.ps.configure(
                max_rows=task.max_rows,
                variations_per_field=task.variations_per_field,
                max_variations_per_row=task.max_variations_per_row,
                random_seed=task.random_seed,
                api_platform=task.api_platform,
                model_name=task.model_name
            )
            print("\n4. Generating prompt variations...")
            variations = task.ps.generate(verbose=True)

            # Display results
            print(f"\n✅ Generated {len(variations)} variations")

            # Show a few examples
            print("\n5. Sample variations:")
            for i, var in enumerate(variations[:3]):
                print(f"\nVariation {i + 1}:")
                print("-" * 50)
                prompt = var.get('prompt', 'No prompt found')
                if len(prompt) > 500:
                    prompt = prompt[:500] + "..."
                print(prompt)
                print("-" * 50)

            # Export results using the correct path
            print(f"\n6. Exporting results to {output_file}...")
            task.ps.export(str(output_file), format="json")
            print("✅ Export completed!")

            # Show final statistics
            print("\n7. Final statistics:")
            task.ps.info()

            generated_files.append(str(output_file))
            print(f"✅ Completed {display_name}: {output_file}")
            
        except Exception as e:
            print(f"❌ Error processing {display_name}: {e}")
            continue
    
    print(f"\n🎉 All language pairs completed! Generated {len(generated_files)} files:")
    for file in generated_files:
        print(f"  📄 {file}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate translation prompt variations")
    parser.add_argument("--language_pair", help="Specific language pair to process (e.g., cs-en, de-en)", default=None)
    parser.add_argument("--all", action="store_true", help="Process all language pairs", default=True)
    parser.add_argument("--rows", type=int, help="Number of rows to process", default=DEFAULT_MAX_ROWS)
    parser.add_argument("--variations", type=int, help="Number of variations per row", default=DEFAULT_MAX_VARIATIONS_PER_ROW)
    parser.add_argument("--variations_per_field", type=int, default=DEFAULT_VARIATIONS_PER_FIELD)
    parser.add_argument("--api_platform", type=str, default=DEFAULT_PLATFORM)
    parser.add_argument("--model_name", type=str, default=DEFAULT_MODEL_NAME)
    parser.add_argument("--random_seed", type=int, help="Random seed for generation", default=DEFAULT_RANDOM_SEED)
    args = parser.parse_args()
    
    if args.all:
        generate_all_language_pairs(
            variations_per_field=args.variations_per_field,
            api_platform=args.api_platform,
            model_name=args.model_name,
            max_rows=args.rows,
            max_variations_per_row=args.variations,
            random_seed=args.random_seed
        )
    elif args.language_pair:
        task = TranslationTask(
            language_pair=args.language_pair,
            variations_per_field=args.variations_per_field,
            api_platform=args.api_platform,
            model_name=args.model_name,
            max_rows=args.rows,
            max_variations_per_row=args.variations,
            random_seed=args.random_seed
        )
        if args.rows is not None or args.variations is not None:
            task.override_config(rows=args.rows, variations=args.variations)
        task.generate()
    else:
        print("Please specify either --language_pair <pair> or --all")
        print("\nAvailable language pairs:")
        try:
            pairs = get_available_language_pairs()
            for pair in pairs:
                display_name = TranslationTask._get_display_name(pair)
                print(f"  - {display_name} ({pair})")
        except Exception as e:
            print(f"Error getting language pairs: {e}") 