import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import (
    AutoTokenizer, 
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    get_scheduler
)
from accelerate import Accelerator
from datasets import Dataset
from tqdm.auto import tqdm
import evaluate
import json
from pathlib import Path
from dataclasses import dataclass
from typing import Dict, List, Optional
import logging
import argparse
import sys
import yaml

logger = logging.getLogger(__name__)

@dataclass
class SimpleConfig:
    """Simple configuration for accelerate training"""
    # Model settings
    model_name: str = "bert-base-uncased"
    max_length: int = 512
    
    # Training settings
    num_epochs: int = 3
    batch_size: int = 16
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    
    # Scheduler settings
    lr_scheduler_type: str = "linear"
    warmup_ratio: float = 0.1
    
    # Paths
    data_dir: str = "./data/classification"
    output_dir: str = "./results"


class AccelerateTrainer:
    """Simple trainer using Accelerate for distributed training"""
    
    def __init__(self, config: SimpleConfig):
        self.config = config
        
        # Initialize accelerator
        self.accelerator = Accelerator()
        
        # Setup logging only on main process
        if self.accelerator.is_main_process:
            logging.basicConfig(
                format='%(asctime)s - %(levelname)s - %(message)s',
                level=logging.INFO
            )
        
        self.tokenizer = None
        self.model = None
        self.label_to_id = {}
        self.id_to_label = {}
        self.num_labels = 0
    
    def load_data(self) -> Dict[str, List[Dict]]:
        """Load data from JSONL files"""
        data_path = Path(self.config.data_dir)
        splits = {}
        
        for split_name in ["train", "validation", "test"]:
            split_file = data_path / f"{split_name}.jsonl"
            if split_file.exists():
                split_data = []
                with open(split_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        if line.strip():
                            split_data.append(json.loads(line))
                splits[split_name] = split_data
                
                if self.accelerator.is_main_process:
                    logger.info(f"Loaded {len(split_data)} samples from {split_name}")
        
        return splits
    
    def setup_labels(self, train_data: List[Dict]):
        """Setup label mappings"""
        labels = set()
        for item in train_data:
            labels.add(str(item["label"]))
        
        sorted_labels = sorted(list(labels))
        self.label_to_id = {label: idx for idx, label in enumerate(sorted_labels)}
        self.id_to_label = {idx: label for label, idx in self.label_to_id.items()}
        self.num_labels = len(sorted_labels)
        
        if self.accelerator.is_main_process:
            logger.info(f"Found {self.num_labels} labels: {sorted_labels}")
    
    def create_dataset(self, data: List[Dict]) -> Dataset:
        """Create tokenized dataset"""
        texts = []
        labels = []
        
        for item in data:
            text = item["text"] if "text" in item else item["input"]
            label = item["label"]
            
            texts.append(str(text))
            # Convert label to ID
            if isinstance(label, str):
                label_id = self.label_to_id.get(label, 0)
            else:
                label_id = int(label)
            labels.append(label_id)
        
        # Create dataset
        dataset = Dataset.from_dict({
            "text": texts,
            "labels": labels
        })
        
        # Tokenize
        def tokenize_function(examples):
            return self.tokenizer(
                examples["text"],
                truncation=True,
                padding="max_length",
                max_length=self.config.max_length
            )
        
        tokenized_dataset = dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=["text"]
        )
        
        return tokenized_dataset
    
    def compute_metrics(self, predictions, labels):
        """Compute accuracy and F1"""
        predictions = predictions.argmax(axis=-1)
        
        # Gather predictions from all processes
        all_predictions = self.accelerator.gather_for_metrics(predictions)
        all_labels = self.accelerator.gather_for_metrics(labels)
        
        if self.accelerator.is_main_process:
            # Only compute metrics on main process
            metric = evaluate.load("glue", "mrpc")  # Using MRPC as example
            results = metric.compute(
                predictions=all_predictions.cpu().numpy(),
                references=all_labels.cpu().numpy()
            )
            return results
        return {}
    
    def train(self):
        """Main training function"""
        if self.accelerator.is_main_process:
            logger.info("=== Starting Accelerate Training ===")
        
        # Load data
        splits_data = self.load_data()
        if "train" not in splits_data:
            raise ValueError("No training data found!")
        
        # Setup labels
        self.setup_labels(splits_data["train"])
        
        # Initialize tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.config.model_name,
            num_labels=self.num_labels,
            id2label=self.id_to_label,
            label2id=self.label_to_id
        )
        
        # Create datasets
        train_dataset = self.create_dataset(splits_data["train"])
        eval_dataset = None
        if "validation" in splits_data:
            eval_dataset = self.create_dataset(splits_data["validation"])
        
        # Create data loaders
        data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
        
        train_dataloader = DataLoader(
            train_dataset,
            shuffle=True,
            batch_size=self.config.batch_size,
            collate_fn=data_collator
        )
        
        eval_dataloader = None
        if eval_dataset:
            eval_dataloader = DataLoader(
                eval_dataset,
                batch_size=self.config.batch_size,
                collate_fn=data_collator
            )
        
        # Setup optimizer and scheduler
        optimizer = AdamW(
            self.model.parameters(),
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay
        )
        
        num_training_steps = self.config.num_epochs * len(train_dataloader)
        num_warmup_steps = int(num_training_steps * self.config.warmup_ratio)
        
        lr_scheduler = get_scheduler(
            self.config.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=num_training_steps
        )
        
        # Prepare everything with accelerator
        self.model, optimizer, train_dataloader, lr_scheduler = self.accelerator.prepare(
            self.model, optimizer, train_dataloader, lr_scheduler
        )
        
        if eval_dataloader:
            eval_dataloader = self.accelerator.prepare(eval_dataloader)
        
        # Training loop
        if self.accelerator.is_main_process:
            progress_bar = tqdm(range(num_training_steps))
            logger.info(f"Training steps: {num_training_steps}")
        
        self.model.train()
        
        for epoch in range(self.config.num_epochs):
            if self.accelerator.is_main_process:
                logger.info(f"Epoch {epoch + 1}/{self.config.num_epochs}")
            
            for step, batch in enumerate(train_dataloader):
                outputs = self.model(**batch)
                loss = outputs.loss
                
                # Backward pass
                self.accelerator.backward(loss)
                
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
                
                if self.accelerator.is_main_process:
                    progress_bar.update(1)
                    
                    # Log every 100 steps
                    if step % 100 == 0:
                        logger.info(f"Step {step}, Loss: {loss.item():.4f}")
            
            # Evaluation at end of each epoch
            if eval_dataloader:
                self.evaluate(eval_dataloader, epoch)
        
        # Save model
        self.save_model()
        
        if self.accelerator.is_main_process:
            logger.info("=== Training Completed ===")
    
    def evaluate(self, eval_dataloader, epoch):
        """Evaluation function"""
        self.model.eval()
        
        all_predictions = []
        all_labels = []
        
        for batch in eval_dataloader:
            with torch.no_grad():
                outputs = self.model(**batch)
            
            predictions = outputs.logits
            labels = batch["labels"]
            
            # Gather from all processes
            predictions = self.accelerator.gather_for_metrics(predictions)
            labels = self.accelerator.gather_for_metrics(labels)
            
            all_predictions.append(predictions.cpu())
            all_labels.append(labels.cpu())
        
        if self.accelerator.is_main_process and all_predictions:
            all_predictions = torch.cat(all_predictions)
            all_labels = torch.cat(all_labels)
            
            predictions_np = all_predictions.argmax(dim=-1).numpy()
            labels_np = all_labels.numpy()
            
            # Simple accuracy calculation
            accuracy = (predictions_np == labels_np).mean()
            logger.info(f"Epoch {epoch + 1} - Validation Accuracy: {accuracy:.4f}")
        
        self.model.train()
    
    def save_model(self):
        """Save model and tokenizer"""
        if self.accelerator.is_main_process:
            output_path = Path(self.config.output_dir)
            output_path.mkdir(parents=True, exist_ok=True)
            
            # Save model using accelerator
            unwrapped_model = self.accelerator.unwrap_model(self.model)
            unwrapped_model.save_pretrained(output_path)
            self.tokenizer.save_pretrained(output_path)
            
            # Save label info
            label_info = {
                "label_to_id": self.label_to_id,
                "id_to_label": self.id_to_label,
                "num_labels": self.num_labels
            }
            
            with open(output_path / "label_info.json", 'w') as f:
                json.dump(label_info, f, indent=2)
            
            logger.info(f"Model saved to {output_path}")


def main():
    """Main function with YAML configuration support"""
    
    parser = argparse.ArgumentParser(description="Accelerate Training Pipeline")
    
    # YAML configuration
    parser.add_argument("--config", type=str, help="Path to YAML configuration file")
    
    # Model settings
    parser.add_argument("--model-name", type=str, help="Model name from HuggingFace Hub")
    parser.add_argument("--max-length", type=int, help="Maximum sequence length for tokenization")
    
    # Training settings
    parser.add_argument("--num-epochs", type=int, help="Number of training epochs")
    parser.add_argument("--batch-size", type=int, help="Training batch size")
    parser.add_argument("--learning-rate", type=float, help="Learning rate")
    parser.add_argument("--weight-decay", type=float, help="Weight decay for optimizer")
    
    # Scheduler settings
    parser.add_argument("--lr-scheduler-type", choices=["linear", "cosine", "polynomial"], help="Learning rate scheduler type")
    parser.add_argument("--warmup-ratio", type=float, help="Warmup ratio for scheduler")
    
    # Paths
    parser.add_argument("--data-dir", type=str, help="Directory containing train/validation/test JSONL files")
    parser.add_argument("--output-dir", type=str, help="Output directory for saved model")
    
    # Logging
    parser.add_argument("--log-level", choices=["DEBUG", "INFO", "WARNING", "ERROR"], default="INFO", help="Logging level")
    
    args = parser.parse_args()
    
    # Set up logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # Load configuration
    config_dict = {}
    
    # Load YAML config if provided
    if args.config:
        try:
            with open(args.config, 'r', encoding='utf-8') as f:
                config_dict = yaml.safe_load(f)
            logger.info(f"Loaded YAML configuration from: {args.config}")
        except Exception as e:
            logger.error(f"Error loading YAML config: {e}")
            sys.exit(1)
    
    # Override YAML config with CLI arguments
    cli_overrides = {}
    if args.model_name:
        cli_overrides['model_name'] = args.model_name
    if args.max_length:
        cli_overrides['max_length'] = args.max_length
    if args.num_epochs:
        cli_overrides['num_epochs'] = args.num_epochs
    if args.batch_size:
        cli_overrides['batch_size'] = args.batch_size
    if args.learning_rate:
        cli_overrides['learning_rate'] = args.learning_rate
    if args.weight_decay:
        cli_overrides['weight_decay'] = args.weight_decay
    if args.lr_scheduler_type:
        cli_overrides['lr_scheduler_type'] = args.lr_scheduler_type
    if args.warmup_ratio:
        cli_overrides['warmup_ratio'] = args.warmup_ratio
    if args.data_dir:
        cli_overrides['data_dir'] = args.data_dir
    if args.output_dir:
        cli_overrides['output_dir'] = args.output_dir
    
    # Merge configurations
    for key, value in cli_overrides.items():
        if key in config_dict:
            logger.info(f"Overriding YAML config '{key}' with CLI value: {value}")
        config_dict[key] = value
    
    # Create configuration object
    config = SimpleConfig(
        model_name=config_dict.get('model_name', 'bert-base-uncased'),
        max_length=config_dict.get('max_length', 512),
        num_epochs=config_dict.get('num_epochs', 3),
        batch_size=config_dict.get('batch_size', 16),
        learning_rate=config_dict.get('learning_rate', 2e-5),
        weight_decay=config_dict.get('weight_decay', 0.01),
        lr_scheduler_type=config_dict.get('lr_scheduler_type', 'linear'),
        warmup_ratio=config_dict.get('warmup_ratio', 0.1),
        data_dir=config_dict.get('data_dir', './data/classification'),
        output_dir=config_dict.get('output_dir', './results')
    )
    
    # Validate data directory
    data_path = Path(config.data_dir)
    if not data_path.exists():
        print(f"❌ Data directory not found: {data_path}")
        print("Please ensure the data directory exists and contains train.jsonl, validation.jsonl, and test.jsonl files")
        sys.exit(1)
    
    # Check for required data files
    required_files = ["train.jsonl"]
    missing_files = []
    for file_name in required_files:
        if not (data_path / file_name).exists():
            missing_files.append(file_name)
    
    if missing_files:
        print(f"❌ Missing required data files: {missing_files}")
        print(f"Please ensure these files exist in: {data_path}")
        sys.exit(1)
    
    # Initialize and run training
    try:
        print(f"Starting training with model: {config.model_name}")
        print(f"Data directory: {config.data_dir}")
        print(f"Output directory: {config.output_dir}")
        print(f"Training for {config.num_epochs} epochs with batch size {config.batch_size}")
        if args.config:
            print(f"Using YAML configuration: {args.config}")
        print()
        
        trainer = AccelerateTrainer(config)
        trainer.train()
        
        print(f"✅ Training completed successfully!")
        print(f"Model saved to: {config.output_dir}")
        
    except Exception as e:
        print(f"❌ Error during training: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()