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

import time
from typing import Optional, Union, Any

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 ..datasets.registry import get_builtin_defaults
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,
        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.generate_kwargs = generate_kwargs or {}
        self.capture_anchor_on_train_begin = capture_anchor_on_train_begin

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

        for d in datasets:
            if isinstance(d, str):
                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

        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."""
        if self.capture_anchor_on_train_begin and not self._anchor_captured:
            if model is not None and tokenizer is not None:
                self.capture_anchor(model, 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

        # 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:
                    self._run_evaluation(model, tokenizer, step, epoch)

    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...")

        # 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
        for ds_config in self.datasets:
            generations = []

            with torch.no_grad():
                for prompt_data in ds_config._prompts:
                    try:
                        start_time = time.time()
                        output = self._generate(model, tokenizer, prompt_data["prompt"])
                        gen_time_ms = int((time.time() - start_time) * 1000)

                        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": gen_time_ms,
                        })
                    except Exception as e:
                        logger.error(f"Error generating anchor for prompt {prompt_data['id']}: {e}")

            # 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}")
        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
        for ds_config in self.datasets:
            generations = []

            with torch.no_grad():
                for prompt_data in ds_config._prompts:
                    try:
                        start_time = time.time()
                        output = self._generate(model, tokenizer, prompt_data["prompt"])
                        gen_time_ms = int((time.time() - start_time) * 1000)

                        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": gen_time_ms,
                        })
                    except Exception as e:
                        logger.error(f"Error generating for prompt {prompt_data['id']}: {e}")

            # 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)
