"""Hugging Face Trainer callback for TrainTrack evaluation integration."""

import time
from typing import Optional, Union, Any
import logging

import torch

# Try to import TrainerCallback, fallback to object if transformers not installed
try:
    from transformers import TrainerCallback
except ImportError:
    TrainerCallback = object

from .sender import TrainTrackSender
from .hook import DatasetConfig, MetricConfig, CategoryConfig
from .ddp import is_rank0, safe_barrier, is_distributed, get_rank
from ..datasets import DatasetLoader
from ._inference_utils import (
    detect_optimal_batch_size,
    try_compile_model,
    get_optimized_generation_config,
)
from ..datasets.registry import get_builtin_defaults
from ..datasets.categories import list_categories
from ..utils import get_logger

logger = get_logger(__name__)

class TrainTrackCallback(TrainerCallback):
    """
    Hugging Face Trainer callback for TrainTrack evaluation.

    This callback integrates with the `transformers.Trainer` to:
    1. Capture anchor outputs before training starts
    2. Run evaluations every N steps during training
    3. Send all outputs to the TrainTrack server for judging

    Works with:
    - Standard HuggingFace models
    - PEFT/LoRA wrapped models
    - Any model that exposes a `.generate()` method
    - DDP/multi-GPU training (only rank 0 evaluates and sends)

    Usage:
        from traintrack.client import TrainTrackCallback, DatasetConfig
        from transformers import Trainer

        callback = TrainTrackCallback(
            run_name="my-lora-finetune",
            datasets=[
                DatasetConfig(name="reasoning", max_samples=10),
            ],
            eval_every_steps=50,
        )

        # Capture anchor BEFORE training
        callback.capture_anchor(model, tokenizer)

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_data,
            callbacks=[callback],
        )
        trainer.train()
    """

    def __init__(
        self,
        run_name: str,
        datasets: list[Union[str, DatasetConfig]],
        eval_every_steps: int = 100,
        max_new_tokens: int = 256,
        batch_size: Optional[int] = None,  # Auto-detect if None
        use_torch_compile: bool = False,    # Enable torch.compile speedup
        generate_kwargs: Optional[dict] = None,
        server_url: Optional[str] = None,
        api_key: Optional[str] = None,
        capture_anchor_on_train_begin: bool = False,
    ):
        """
        Initialize the TrainTrack callback.

        Args:
            run_name: Name for this training run
            datasets: List of dataset names or DatasetConfig objects
            eval_every_steps: Run evaluation every N steps
            max_new_tokens: Max tokens to generate per prompt
            generate_kwargs: Additional kwargs for model.generate()
            server_url: TrainTrack server URL (optional, uses env var)
            api_key: TrainTrack API key (optional, uses env var)
            capture_anchor_on_train_begin: If True, auto-capture anchor in on_train_begin
        """
        self.run_name = run_name
        self.eval_every_steps = eval_every_steps
        self.max_new_tokens = max_new_tokens
        self.batch_size_config = batch_size  # Store config, actual batch_size set on first use
        self.use_torch_compile = use_torch_compile
        self._model_compiled = False
        self._batch_size = None  # Will be set dynamically
        
        # Merge user kwargs with optimized defaults (user kwargs take precedence)
        optimized_kwargs = get_optimized_generation_config()
        optimized_kwargs.update(generate_kwargs or {})
        self.generate_kwargs = optimized_kwargs
        
        self.capture_anchor_on_train_begin = capture_anchor_on_train_begin

        # Prepare datasets
        self.loader = DatasetLoader()
        self.datasets: list[Union[DatasetConfig, CategoryConfig]] = []

        valid_categories = list_categories()
        for d in datasets:
            if isinstance(d, str):
                if d in valid_categories:
                    self.datasets.append(CategoryConfig(category=d))
                else:
                    self.datasets.append(DatasetConfig(name=d))
            else:
                self.datasets.append(d)

        # Load prompts for all datasets/categories
        total_samples = 0
        for ds_config in self.datasets:
            if isinstance(ds_config, CategoryConfig):
                # Load from category sources
                ds_config._prompts = self.loader.load_category(
                    ds_config.category,
                    max_samples=ds_config.max_samples,
                    seed=ds_config.sample_seed
                )
            else:
                # Load from dataset file
                ds_config._prompts = self.loader.load(
                    ds_config.name,
                    max_samples=ds_config.max_samples,
                    seed=ds_config.sample_seed
                )
            total_samples += len(ds_config._prompts)

        # Sender (only initialized on rank 0)
        self._is_rank0 = is_rank0()
        self.sender = None
        if self._is_rank0:
            self.sender = TrainTrackSender(server_url=server_url, api_key=api_key)
            self.sender.verify_connection()

        # State
        self._last_eval_step = -1
        self._anchor_captured = False
        self._model = None  # Will be set in on_train_begin
        self._tokenizer = None  # Will be set in on_train_begin

        if self._is_rank0:
            logger.info(
                f"TrainTrackCallback initialized: run={run_name}, "
                f"datasets={len(self.datasets)}, total_samples={total_samples}"
            )
        
        if is_distributed():
            logger.info(f"DDP mode detected (rank {get_rank()}), only rank 0 will evaluate")

    # =========================================================================
    # Hugging Face TrainerCallback interface methods
    # =========================================================================

    def on_train_begin(self, args, state, control, model=None, tokenizer=None, **kwargs):
        """Called at the beginning of training."""
        # Store references for use in on_step_end (HF Trainer may not pass them every time)
        if model is not None:
            self._model = model
        if tokenizer is not None:
            self._tokenizer = tokenizer
        
        if self.capture_anchor_on_train_begin and not self._anchor_captured:
            if self._model is not None and self._tokenizer is not None:
                self.capture_anchor(self._model, self._tokenizer)
            else:
                logger.warning(
                    "capture_anchor_on_train_begin=True but model/tokenizer not available. "
                    "Call capture_anchor() manually before training."
                )

    def on_step_end(self, args, state, control, model=None, tokenizer=None, **kwargs):
        """Called at the end of each training step."""
        step = state.global_step
        epoch = state.epoch
        
        # Use stored refs if not provided (SFTTrainer doesn't pass tokenizer to on_step_end)
        model = model or self._model
        tokenizer = tokenizer or self._tokenizer

        # Check if we should evaluate
        if step > 0 and step % self.eval_every_steps == 0:
            if step != self._last_eval_step:
                self._last_eval_step = step
                if model is not None and tokenizer is not None:
                    logger.info(f"Triggering TrainTrack evaluation at step {step}")
                    self._run_evaluation(model, tokenizer, step, epoch)
                else:
                    logger.warning(f"Cannot run evaluation at step {step}: model or tokenizer not available")

    def on_train_end(self, args, state, control, **kwargs):
        """Called at the end of training."""
        logger.info(f"Training complete for run: {self.run_name}")

    # =========================================================================
    # Core evaluation methods
    # =========================================================================

    def capture_anchor(self, model, tokenizer):
        """
        Capture anchor outputs from the model in its current state.
        Call this BEFORE training starts to establish a baseline for pairwise comparisons.

        Args:
            model: The model (can be a PeftModel or standard HF model)
            tokenizer: The tokenizer
        """
        logger.info("Capturing anchor outputs...")

        # Lazy initialization: compile model and detect batch size on first use
        if not self._model_compiled and self.use_torch_compile:
            model = try_compile_model(model, enable=self.use_torch_compile)
            self._model_compiled = True
        
        if self._batch_size is None:
            if self.batch_size_config is None:
                self._batch_size = detect_optimal_batch_size(model)
            else:
                self._batch_size = self.batch_size_config
                logger.info(f"Using user-configured batch_size={self._batch_size}")


        # Synchronize all ranks before anchor capture
        safe_barrier() 
        
        # Only rank 0 performs the actual capture
        if not self._is_rank0:
            self._anchor_captured = True
            safe_barrier()  # Wait for rank 0 to finish
            return

        checkpoint_tag = "anchor"
        step = 0

        model.eval()

        # Prepare judge config
        judge_config = {"datasets": {}}
        for ds in self.datasets:
            ds_conf = {
                "modes": ds.judge_modes,
                "metrics": []
            }
            for m in ds.metrics:
                if isinstance(m, MetricConfig):
                    ds_conf["metrics"].append({"name": m.name, "rubric": m.rubric})
                else:
                    ds_conf["metrics"].append({"name": m})
            judge_config["datasets"][ds._dataset_name_clean] = ds_conf

        # Process each dataset with batched generation
        for ds_config in self.datasets:
            generations = []

            with torch.no_grad():
                # Extract prompts for batching
                prompts = [p["prompt"] for p in ds_config._prompts]
                
                # Generate in batches
                start_time = time.time()
                outputs = self._generate_batch(model, tokenizer, prompts, batch_size=self._batch_size)
                total_time_ms = int((time.time() - start_time) * 1000)
                avg_time_ms = total_time_ms // len(outputs) if outputs else 0
                
                # Package results
                for prompt_data, output in zip(ds_config._prompts, outputs):
                    generations.append({
                        "prompt": {
                            "id": prompt_data["id"],
                            "dataset_name": ds_config._dataset_name_clean,
                            "prompt": prompt_data["prompt"],
                            "category": prompt_data.get("category"),
                            "reference": prompt_data.get("reference"),
                            "competitor_output": prompt_data.get("competitor_output"),
                            "metadata": prompt_data.get("metadata"),
                        },
                        "output": output,
                        "generation_time_ms": avg_time_ms,
                    })

            # Send to server
            metadata = {"is_anchor": True, "judge_config": judge_config}

            success = self.sender.send_generations(
                run_name=self.run_name,
                checkpoint_tag=checkpoint_tag,
                step=step,
                epoch=0,
                generations=generations,
                metadata=metadata,
            )

            if success:
                logger.info(f"✓ Anchor outputs captured for {ds_config._dataset_name_clean}")
            else:
                logger.error(f"✗ Failed to send anchor outputs for {ds_config._dataset_name_clean}")

        model.train()
        self._anchor_captured = True
        
        # Synchronize all ranks after anchor capture
        safe_barrier()

    def _run_evaluation(self, model, tokenizer, step: int, epoch: Optional[float]):
        """Run evaluation on all datasets and send to server."""
        # Synchronize all ranks before evaluation
        safe_barrier()
        
        # Only rank 0 performs the actual evaluation
        if not self._is_rank0:
            safe_barrier()  # Wait for rank 0 to finish
            return
        
        logger.info(f"Running evaluation at step {step}")
        
        # Ensure batch_size is set (lazy init if not done in capture_anchor)
        if self._batch_size is None:
            if self.batch_size_config is None:
                self._batch_size = detect_optimal_batch_size(model)
            else:
                self._batch_size = self.batch_size_config
                logger.info(f"Using user-configured batch_size={self._batch_size}")
        
        model.eval()

        # Prepare judge config
        judge_config = {"datasets": {}}
        for ds in self.datasets:
            ds_conf = {
                "modes": ds.judge_modes,
                "metrics": []
            }
            for m in ds.metrics:
                if isinstance(m, MetricConfig):
                    ds_conf["metrics"].append({"name": m.name, "rubric": m.rubric})
                else:
                    ds_conf["metrics"].append({"name": m})
            judge_config["datasets"][ds._dataset_name_clean] = ds_conf

        # Process each dataset with batched generation
        for ds_config in self.datasets:
            generations = []

            with torch.no_grad():
                # Extract prompts for batching
                prompts = [p["prompt"] for p in ds_config._prompts]
                
                # Generate in batches
                start_time = time.time()
                outputs = self._generate_batch(model, tokenizer, prompts, batch_size=self._batch_size)
                total_time_ms = int((time.time() - start_time) * 1000)
                avg_time_ms = total_time_ms // len(outputs) if outputs else 0
                
                # Package results
                for prompt_data, output in zip(ds_config._prompts, outputs):
                    generations.append({
                        "prompt": {
                            "id": prompt_data["id"],
                            "dataset_name": ds_config._dataset_name_clean,
                            "prompt": prompt_data["prompt"],
                            "category": prompt_data.get("category"),
                            "reference": prompt_data.get("reference"),
                            "competitor_output": prompt_data.get("competitor_output"),
                            "metadata": prompt_data.get("metadata"),
                        },
                        "output": output,
                        "generation_time_ms": avg_time_ms,
                    })

            # Send to server
            checkpoint_tag = f"step_{step}"
            if epoch is not None:
                checkpoint_tag = f"epoch_{int(epoch)}_step_{step}"

            logger.info(f"Sending {len(generations)} generations for {ds_config._dataset_name_clean}")

            success = self.sender.send_generations(
                run_name=self.run_name,
                checkpoint_tag=checkpoint_tag,
                step=step,
                epoch=int(epoch) if epoch else None,
                generations=generations,
                metadata={"judge_config": judge_config},
            )

            if not success:
                logger.warning(f"Failed to send generations for step {step}")

        model.train()
        
        # Synchronize all ranks after evaluation
        safe_barrier()

    def _generate(self, model, tokenizer, prompt: str) -> str:
        """Generate output for a single prompt."""
        device = next(model.parameters()).device

        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=1024,
        )

        # Handle both BatchEncoding and plain dicts
        if hasattr(inputs, "to"):
            inputs = inputs.to(device)
        else:
            inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}

        outputs = model.generate(
            inputs["input_ids"],
            attention_mask=inputs.get("attention_mask"),
            max_new_tokens=self.max_new_tokens,
            pad_token_id=getattr(tokenizer, "pad_token_id", None) or getattr(tokenizer, "eos_token_id", None),
            **self.generate_kwargs,
        )

        # Decode only new tokens
        input_len = inputs["input_ids"].shape[1]
        new_tokens = outputs[0][input_len:]
        return tokenizer.decode(new_tokens, skip_special_tokens=True)

    def _generate_batch(self, model, tokenizer, prompts: list[str], batch_size: int = 4) -> list[str]:
        """Generate outputs for multiple prompts in batches for faster inference."""
        device = next(model.parameters()).device
        all_outputs = []
        
        # Enforce left padding for generation (critical for decoder-only models)
        old_padding_side = getattr(tokenizer, "padding_side", None)
        if old_padding_side != "left":
            tokenizer.padding_side = "left"
            
        try:
            for i in range(0, len(prompts), batch_size):
                batch_prompts = prompts[i:i + batch_size]
                
                # Tokenize batch with padding
                inputs = tokenizer(
                    batch_prompts,
                    return_tensors="pt",
                    truncation=True,
                    max_length=1024,
                    padding=True,
                )
                
                # Move to device
                if hasattr(inputs, "to"):
                    inputs = inputs.to(device)
                else:
                    inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}
                
                # Generate
                outputs = model.generate(
                    inputs["input_ids"],
                    attention_mask=inputs.get("attention_mask"),
                    max_new_tokens=self.max_new_tokens,
                    pad_token_id=getattr(tokenizer, "pad_token_id", None) or getattr(tokenizer, "eos_token_id", None),
                    **self.generate_kwargs,
                )
                
                # Decode each output (only new tokens)
                for j, output in enumerate(outputs):
                    input_len = inputs["input_ids"][j].shape[0]
                    new_tokens = output[input_len:]
                    decoded = tokenizer.decode(new_tokens, skip_special_tokens=True)
                    if not decoded.strip():
                        print(f"[TrainTrack] WARNING: Empty generation for prompt index {i+j}!")
                    all_outputs.append(decoded)
        finally:
            if old_padding_side is not None:
                tokenizer.padding_side = old_padding_side
        
        return all_outputs

