"""Training hook for TrainTrack evaluation integration."""

import time
from dataclasses import dataclass, field
from typing import Optional, Callable, Any, Union

import torch

from .ddp import DDPContext, is_rank0
from .sender import TrainTrackSender
from ..datasets import DatasetLoader
from ..datasets.registry import get_builtin_defaults
from ..datasets.categories import get_category_config, get_category_rubrics
from ..utils import get_logger

logger = get_logger(__name__)


@dataclass
class MetricConfig:
    """Configuration for a specific metric."""

    name: str
    rubric: Optional[str] = None

    # Predefined rubrics can be referenced by name if supported by server
    # or fully specified here


@dataclass
class DatasetConfig:
    """Configuration for a specific dataset to evaluate."""

    name: str  # Built-in name or path to JSONL/CSV
    max_samples: int = 50
    sample_seed: Optional[int] = 42
    judge_modes: list[str] = field(default_factory=list)
    metrics: list[Union[str, MetricConfig]] = field(default_factory=list)

    # Internal use
    _prompts: list[dict] = field(default_factory=list, repr=False)
    _dataset_name_clean: str = field(default="", repr=False)

    def __post_init__(self):
        if "/" in self.name or "\\" in self.name:
            import os

            self._dataset_name_clean = os.path.splitext(os.path.basename(self.name))[0]
        else:
            self._dataset_name_clean = self.name

        # Apply defaults for built-in datasets if not configured
        if not self.metrics and not self.judge_modes:
            defaults = get_builtin_defaults(self.name)
            if defaults:
                if "metrics" in defaults:
                    self.metrics = defaults["metrics"]
                if "judge_modes" in defaults:
                    self.judge_modes = defaults["judge_modes"]

        # Default fallback if still empty
        if not self.judge_modes:
            self.judge_modes = ["criteria"]


@dataclass
class CategoryConfig:
    """
    Configuration for evaluating on a subject category (e.g., physics, law, math).
    
    Categories aggregate prompts from multiple sources (gpqa_diamond, mmlu_pro) and
    automatically apply appropriate domain-specific metrics and rubrics.
    
    Example:
        CategoryConfig(
            category="physics",
            max_samples=50,
        )
        # This will:
        # - Load physics questions from gpqa_diamond AND mmlu_pro
        # - Apply physics-specific metrics: correctness, reasoning_quality, scientific_accuracy
        # - Use both 'criteria' and 'pairwise_anchor' judge modes
    """
    category: str  # Category name: physics, chemistry, math, law, etc.
    max_samples: int = 100
    sample_seed: Optional[int] = 42
    # Override defaults if needed
    judge_modes: Optional[list[str]] = None
    metrics: Optional[list[Union[str, MetricConfig]]] = None
    
    # Internal use
    _prompts: list[dict] = field(default_factory=list, repr=False)
    _config: dict = field(default_factory=dict, repr=False)
    _rubrics: dict = field(default_factory=dict, repr=False)
    
    def __post_init__(self):
        # Load category configuration
        self._config = get_category_config(self.category)
        self._rubrics = get_category_rubrics(self.category)
        
        # Apply category defaults if not overridden
        if self.judge_modes is None:
            self.judge_modes = self._config.get("judge_modes", ["criteria", "pairwise_anchor"])
        if self.metrics is None:
            # Convert metrics to MetricConfig with rubrics
            self.metrics = [
                MetricConfig(name=m, rubric=self._rubrics.get(m))
                for m in self._config.get("metrics", [])
            ]
    
    @property
    def _dataset_name_clean(self) -> str:
        """For compatibility with DatasetConfig interface."""
        return self.category


class TrainTrackHook:
    """
    Hook for integrating TrainTrack evaluation into training loops.

    This hook:
    1. Triggers evaluation every N steps or epochs
    2. Generates model outputs on eval prompts
    3. Sends outputs to the TrainTrack server
    4. Resumes training immediately (doesn't wait for judging)

    Usage:
        # Simple usage (legacy)
        hook = TrainTrackHook(
            model=model,
            tokenizer=tokenizer,
            run_name="my-finetune",
            dataset="reasoning",
        )

        # Advanced usage (multiple datasets with config)
        hook = TrainTrackHook(
            model=model,
            tokenizer=tokenizer,
            run_name="my-finetune",
            datasets=[
                DatasetConfig(
                    name="reasoning",
                    judge_modes=["criteria"],
                    metrics=["reasoning", "helpfulness"]
                ),
                DatasetConfig(
                    name="./custom_eval.jsonl",
                    judge_modes=["pairwise_anchor"],
                    metrics=[
                        MetricConfig(name="politeness", rubric="0=rude, 10=polite")
                    ]
                )
            ]
        )

        for step, batch in enumerate(dataloader):
            # ... training code ...

            hook.step(step, epoch=epoch)
    """

    def __init__(
        self,
        model: torch.nn.Module,
        tokenizer: Any,
        run_name: str,
        # Legacy single-dataset args
        dataset: Optional[str] = None,
        max_samples: int = 50,
        sample_seed: Optional[int] = 42,
        judge_modes: Optional[list[str]] = None,
        # New multi-dataset/category arg
        datasets: Optional[list[Union[DatasetConfig, CategoryConfig]]] = None,
        # Common args
        server_url: Optional[str] = None,
        api_key: Optional[str] = None,
        eval_every_steps: Optional[int] = None,
        eval_every_epochs: Optional[int] = None,
        max_new_tokens: int = 256,
        generate_kwargs: Optional[dict] = None,
    ):
        """
        Initialize the TrainTrack hook.

        Args:
            model: The PyTorch model to evaluate
            tokenizer: Tokenizer with encode/decode methods
            run_name: Name for this training run
            dataset: (Deprecated) Dataset name/path
            max_samples: (Deprecated) Max samples for single dataset
            sample_seed: (Deprecated) Seed for single dataset
            judge_modes: (Deprecated) Judge modes for single dataset
            datasets: List of DatasetConfig or CategoryConfig objects (preferred)
            server_url: TrainTrack server URL (optional)
            api_key: TrainTrack API key
            eval_every_steps: Evaluate every N steps
            eval_every_epochs: Evaluate every N epochs
            max_new_tokens: Max tokens to generate
            generate_kwargs: Additional kwargs for model.generate()
        """
        self.model = model
        self.tokenizer = tokenizer
        self.run_name = run_name
        self.eval_every_steps = eval_every_steps
        self.eval_every_epochs = eval_every_epochs
        self.max_new_tokens = max_new_tokens
        self.generate_kwargs = generate_kwargs or {}
        self.judge_modes = judge_modes

        # At least one trigger must be set
        if eval_every_steps is None and eval_every_epochs is None:
            self.eval_every_steps = 100  # Default to every 100 steps

        # Prepare datasets
        self.loader = DatasetLoader()
        self.datasets = []

        if datasets:
            self.datasets = []
            for d in datasets:
                if isinstance(d, str):
                    self.datasets.append(DatasetConfig(name=d))
                elif isinstance(d, CategoryConfig):
                    # CategoryConfig is used directly
                    self.datasets.append(d)
                else:
                    self.datasets.append(d)
        elif dataset:
            # Backward compatibility
            self.datasets = [
                DatasetConfig(
                    name=dataset,
                    max_samples=max_samples,
                    sample_seed=sample_seed,
                    judge_modes=judge_modes or ["criteria"],
                )
            ]
        else:
            raise ValueError("Must provide either 'datasets' or 'dataset'")
            
        # 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
        self.sender = TrainTrackSender(server_url=server_url, api_key=api_key)

        # Verify connection immediately
        self.sender.verify_connection()

        # State
        self._last_eval_step = -1
        self._last_eval_epoch = -1
        self._step_count = 0

        logger.info(
            f"TrainTrackHook initialized: run={run_name}, "
            f"datasets={len(self.datasets)}, total_samples={total_samples}"
        )

    def step(self, step: int, epoch: Optional[int] = None):
        """
        Call this every training step.

        Args:
            step: Current global step
            epoch: Current epoch (optional)
        """
        self._step_count = step

        should_eval = False

        # Check step-based trigger
        if self.eval_every_steps is not None:
            if step > 0 and step % self.eval_every_steps == 0:
                if step != self._last_eval_step:
                    should_eval = True
                    self._last_eval_step = step

        # Check epoch-based trigger
        if self.eval_every_epochs is not None and epoch is not None:
            if epoch > 0 and epoch % self.eval_every_epochs == 0:
                if epoch != self._last_eval_epoch:
                    should_eval = True
                    self._last_eval_epoch = epoch

        if should_eval:
            self._run_evaluation(step, epoch)

    def _run_evaluation(self, step: int, epoch: Optional[int]):
        """Run evaluation and send to server."""
        with DDPContext() as ctx:
            if not ctx.is_main:
                return

            logger.info(f"Running evaluation at step {step}")
            self.model.eval()

            # Prepare judge config for metadata
            # We construct a map of dataset_name -> 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(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 (non-blocking from training perspective)
                checkpoint_tag = f"step_{step}"
                if epoch is not None:
                    checkpoint_tag = f"epoch_{epoch}_step_{step}"

                # We attach the FULL judge config to every batch to ensure server has it
                # Optimization: could only send once, but stateless is safer
                logger.info(
                    f"Sending {len(generations)} generations for dataset {ds_config._dataset_name_clean}"
                )

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

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

            self.model.train()

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

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

        # Handle both BatchEncoding (with .to()) and plain dicts
        if hasattr(inputs, "to"):
            inputs = inputs.to(device)
        else:
            # Plain dict - move each tensor manually
            inputs = {
                k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()
            }

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

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

    def force_eval(self, step: Optional[int] = None, epoch: Optional[int] = None):
        """Force an immediate evaluation."""
        step = step or self._step_count
        self._run_evaluation(step, epoch)

    def capture_anchor(self):
        """
        Capture anchor outputs from the model in its current state.
        Call this before training starts to establish a baseline.
        """
        logger.info("Capturing anchor outputs...")

        # Determine anchor tag
        checkpoint_tag = "anchor"
        step = 0  # Anchor is technically step 0

        with DDPContext() as ctx:
            if not ctx.is_main:
                return
            self.model.eval()

            # Prepare config (same as in _run_evaluation)
            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(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 and sent for {ds_config._dataset_name_clean}"
                    )
                else:
                    logger.error(
                        f"✗ Failed to send anchor outputs for {ds_config._dataset_name_clean}"
                    )

            self.model.train()
