"""Training orchestrator for UD Parser."""

import json
from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

import torch
from torch import Tensor
from torch.optim import AdamW
from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader
from tqdm import tqdm

from ..config import Config
from ..evaluation.metrics import EvaluationResult
from ..model import UDParser
from ..utils import WandbLogger


class Trainer:
    """Training orchestrator for UD Parser."""

    def __init__(
        self,
        model: UDParser,
        train_dataloader: DataLoader,
        dev_dataloader: DataLoader,
        config: Config,
        evaluator: Any,  # Type hint would create circular import
        logger: WandbLogger,
        device: torch.device,
        test_dataloader: Optional[DataLoader] = None,
        tokenizer: Any = None,
        dev_paths: Any = None,
        test_paths: Any = None,
    ):
        """Initialize trainer.

        Args:
            model: UDParser model.
            train_dataloader: Training data loader.
            dev_dataloader: Development data loader.
            config: Training configuration.
            evaluator: Evaluator instance.
            logger: Wandb logger.
            device: Device for training.
            test_dataloader: Optional test data loader.
            tokenizer: HuggingFace tokenizer (for full-text evaluation).
            dev_paths: SplitPaths with .txt and .conllu paths for dev.
            test_paths: SplitPaths with .txt and .conllu paths for test.
        """
        self.model = model.to(device)
        self.train_dataloader = train_dataloader
        self.dev_dataloader = dev_dataloader
        self.test_dataloader = test_dataloader
        self.config = config
        self.evaluator = evaluator
        self.logger = logger
        self.device = device
        self.tokenizer = tokenizer
        self.dev_paths = dev_paths
        self.test_paths = test_paths

        # Setup optimizer with differential learning rates
        self.optimizer = self._setup_optimizer()
        self.scheduler = self._setup_scheduler()

        # Tracking
        self.best_metric = -1.0
        self.global_step = 0
        self.current_epoch = 0

        # Output directory
        self.output_dir = config.language_output_dir
        self.output_dir.mkdir(parents=True, exist_ok=True)
        (self.output_dir / "checkpoints").mkdir(exist_ok=True)
        (self.output_dir / "results").mkdir(exist_ok=True)

    def _setup_optimizer(self) -> AdamW:
        """Setup AdamW optimizer with differential learning rates."""
        encoder_params = []
        head_params = []

        for name, param in self.model.named_parameters():
            if not param.requires_grad:
                continue
            if "encoder" in name:
                encoder_params.append(param)
            else:
                head_params.append(param)

        optimizer = AdamW(
            [
                {"params": encoder_params, "lr": self.config.training.encoder_lr},
                {"params": head_params, "lr": self.config.training.head_lr},
            ],
            weight_decay=self.config.training.weight_decay,
        )

        print(f"Optimizer setup:")
        print(f"  Encoder params: {len(encoder_params)}, lr={self.config.training.encoder_lr}")
        print(f"  Head params: {len(head_params)}, lr={self.config.training.head_lr}")

        return optimizer

    def _setup_scheduler(self) -> OneCycleLR:
        """Setup learning rate scheduler."""
        accum_steps = self.config.training.gradient_accumulation_steps
        steps_per_epoch = (len(self.train_dataloader) + accum_steps - 1) // accum_steps
        total_steps = steps_per_epoch * self.config.training.epochs

        # Ensure minimum steps for scheduler to work
        if total_steps < 1:
            total_steps = 1
            print("  Warning: Very few training steps, using minimum of 1")

        # Adjust warmup ratio for small datasets (OneCycleLR needs at least some steps)
        warmup_ratio = self.config.training.warmup_ratio
        if total_steps < 10:
            warmup_ratio = 0.0  # Disable warmup for very small datasets
            print(f"  Warning: Disabling warmup due to small dataset ({total_steps} steps)")

        warmup_steps = int(total_steps * warmup_ratio)

        scheduler = OneCycleLR(
            self.optimizer,
            max_lr=[self.config.training.encoder_lr, self.config.training.head_lr],
            total_steps=total_steps,
            pct_start=warmup_ratio,
            anneal_strategy="cos",
        )

        print(f"Scheduler setup:")
        print(f"  Total steps: {total_steps}")
        print(f"  Warmup steps: {warmup_steps}")

        return scheduler

    def train(self) -> dict[str, Any]:
        """Main training loop.

        Returns:
            Dictionary with final metrics.
        """
        print(f"\nStarting training for {self.config.training.epochs} epochs")
        print(f"Training samples: {len(self.train_dataloader.dataset)}")
        print(f"Dev samples: {len(self.dev_dataloader.dataset)}")
        if self.test_dataloader is not None:
            print(f"Test samples: {len(self.test_dataloader.dataset)}")
        print(f"Batch size: {self.config.data.train_batch_size}")
        print(f"Total trainable params: {self.model.get_trainable_params():,}")
        print(f"Best model selection: {self.config.training.metric_for_best} on {self.config.training.best_model_split}")
        print(f"Device: {self.device}", flush=True)

        all_results = []
        fulltext_sources: dict[str, int] = {}  # text_source → count

        for epoch in range(self.config.training.epochs):
            self.current_epoch = epoch

            # Scheduled sampling ratio is updated per-step in _train_epoch
            # (smooth linear ramp 0→1 over all training steps)

            # Training epoch
            train_loss = self._train_epoch()

            # Evaluate on dev set (aligned — gold tokenization)
            print(f"Evaluating on dev set...", flush=True)
            dev_result = self.evaluator.evaluate(
                self.model, self.dev_dataloader, self.device,
                show_progress=True,
            )

            # Full-text evaluation on dev
            dev_fulltext_result = None
            if self.dev_paths and self.dev_paths.conllu_path:
                print(f"Full-text evaluation on dev set...", flush=True)
                dev_fulltext_result = self.evaluator.evaluate_full_text(
                    self.model,
                    self.dev_paths.text_path if self.dev_paths.text_path else None,
                    self.dev_paths.conllu_path,
                    self.tokenizer, self.config.data.max_length,
                    self.config.data.eval_batch_size, self.device,
                    language=self.config.data.language,
                    combo_model_id=self.config.repo_id,
                    show_progress=True,
                )
                if dev_fulltext_result.text_source:
                    fulltext_sources[dev_fulltext_result.text_source] = (
                        fulltext_sources.get(dev_fulltext_result.text_source, 0) + 1
                    )

            # Evaluate on train subset if configured
            train_result = None
            if self.config.training.eval_on_train:
                n_samples = self.config.training.train_eval_samples
                print(f"Evaluating on train set (max_samples={n_samples})...", flush=True)
                train_result = self.evaluator.evaluate(
                    self.model,
                    self.train_dataloader,
                    self.device,
                    max_samples=n_samples,
                    show_progress=True,
                )

            # Evaluate on test set if available
            test_result = None
            test_fulltext_result = None
            if self.test_dataloader is not None:
                print(f"Evaluating on test set...", flush=True)
                test_result = self.evaluator.evaluate(
                    self.model, self.test_dataloader, self.device,
                    show_progress=True,
                )

                # Full-text evaluation on test
                if self.test_paths and self.test_paths.conllu_path:
                    print(f"Full-text evaluation on test set...", flush=True)
                    test_fulltext_result = self.evaluator.evaluate_full_text(
                        self.model,
                        self.test_paths.text_path if self.test_paths.text_path else None,
                        self.test_paths.conllu_path,
                        self.tokenizer, self.config.data.max_length,
                        self.config.data.eval_batch_size, self.device,
                        language=self.config.data.language,
                        combo_model_id=self.config.repo_id,
                        show_progress=True,
                    )
                    if test_fulltext_result.text_source:
                        fulltext_sources[test_fulltext_result.text_source] = (
                            fulltext_sources.get(test_fulltext_result.text_source, 0) + 1
                        )

            # Log metrics to wandb
            self._log_epoch_metrics(
                epoch, train_loss,
                dev_result, train_result, test_result,
                dev_fulltext_result, test_fulltext_result,
            )

            # Select metrics for best model based on configured split
            best_split = self.config.training.best_model_split
            if best_split == "test" and test_result:
                best_metrics = test_result.to_dict()
            elif best_split == "train" and train_result:
                best_metrics = train_result.to_dict()
            else:
                best_metrics = dev_result.to_dict()

            # Save results to JSON
            epoch_results = {
                "epoch": epoch + 1,
                "train_loss": train_loss,
                "dev_metrics": self._result_to_full_metrics(dev_result),
                "train_metrics": self._result_to_full_metrics(train_result),
                "test_metrics": self._result_to_full_metrics(test_result),
                "dev_fulltext_metrics": self._result_to_full_metrics(dev_fulltext_result),
                "test_fulltext_metrics": self._result_to_full_metrics(test_fulltext_result),
            }
            all_results.append(epoch_results)
            self._save_results(epoch_results, "epoch", epoch + 1)

            # Save checkpoint (and results_best.json when improved)
            self._save_checkpoint(epoch, best_metrics, epoch_results)

            # Print summary
            self._print_epoch_summary(
                epoch, train_loss,
                dev_result, train_result, test_result,
                dev_fulltext_result, test_fulltext_result,
            )

        # Save final results
        final_results = {
            "epochs": all_results,
            "best_metric": self.best_metric,
            "config": {
                "data": asdict(self.config.data),
                "model": {
                    "base_model": self.config.model.base_model,
                    "use_lora": self.config.model.use_lora,
                    "lora_r": self.config.model.lora_r,
                },
                "training": asdict(self.config.training),
            },
        }
        self._save_results(final_results, "final")

        self.logger.summary({"best_" + self.config.training.metric_for_best: self.best_metric})

        # Print full-text evaluation source summary
        if fulltext_sources:
            source_labels = {
                "txt_file": ".txt file",
                "text_metadata": "# text = metadata",
                "concatenated_tokens": "concatenated tokens (no .txt, no # text =)",
            }
            print("\nFull-text evaluation text sources:")
            for source, count in fulltext_sources.items():
                print(f"  {source_labels.get(source, source)}: {count}x")

        return final_results

    def _train_epoch(self) -> float:
        """Train for one epoch.

        Returns:
            Average training loss.
        """
        self.model.train()
        total_loss = 0.0
        num_batches = 0

        progress = tqdm(
            self.train_dataloader,
            desc=f"Epoch {self.current_epoch + 1}",
        )

        accum_steps = self.config.training.gradient_accumulation_steps
        steps_per_epoch = len(self.train_dataloader)
        total_training_steps = steps_per_epoch * self.config.training.epochs

        for batch_idx, batch in enumerate(progress):
            # Update scheduled sampling ratio per-step (smooth linear ramp 0→1)
            current_step = self.current_epoch * steps_per_epoch + batch_idx
            ss_ratio = current_step / max(1, total_training_steps - 1)
            self.model.lemma_head.scheduled_sampling_ratio = ss_ratio

            # Move batch to device
            batch = self._move_batch_to_device(batch)

            # Forward pass with mixed precision
            with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.device.type == "cuda"):
                outputs = self.model(**batch)
            loss = outputs["loss"] / accum_steps

            # Backward pass
            loss.backward()

            # Track (use unscaled loss for logging)
            total_loss += outputs["loss"].item()
            num_batches += 1

            # Update weights every accum_steps or at last batch
            if (batch_idx + 1) % accum_steps == 0 or (batch_idx + 1) == len(self.train_dataloader):
                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(),
                    self.config.training.gradient_clip,
                )

                # Update
                self.optimizer.step()
                self.optimizer.zero_grad()
                # Guard against scheduler stepping past total_steps
                if self.global_step < self.scheduler.total_steps:
                    self.scheduler.step()
                self.global_step += 1

                # Free cached memory on MPS/CUDA
                if self.device.type == "mps":
                    torch.mps.empty_cache()
                elif self.device.type == "cuda":
                    torch.cuda.empty_cache()

            # Log to wandb
            log_dict = {
                "global_step": self.global_step,
                "train/loss": outputs["loss"].item(),
                "train/morpho_loss": outputs.get("morpho_loss", torch.tensor(0)).item(),
                "train/upos_loss": outputs.get("upos_loss", torch.tensor(0)).item(),
                "train/xpos_loss": outputs.get("xpos_loss", torch.tensor(0)).item(),
                "train/feats_loss": outputs.get("feats_loss", torch.tensor(0)).item(),
                "train/arc_loss": outputs.get("arc_loss", torch.tensor(0)).item(),
                "train/rel_loss": outputs.get("rel_loss", torch.tensor(0)).item(),
                "train/lemma_loss": outputs.get("lemma_loss", torch.tensor(0)).item(),
                "train/lr_encoder": self.scheduler.get_last_lr()[0],
                "train/lr_head": self.scheduler.get_last_lr()[1],
                "train/ss_ratio": ss_ratio,
            }
            # Log uncertainty weights if available
            if "uncertainty_weights" in outputs:
                for task_name, weight_val in outputs["uncertainty_weights"].items():
                    log_dict[f"train/uw_{task_name}"] = weight_val
            self.logger.log(log_dict)

            progress.set_postfix(loss=outputs["loss"].item())

        return total_loss / num_batches

    def _move_batch_to_device(self, batch: dict[str, Any]) -> dict[str, Any]:
        """Move batch tensors to device."""
        result = {}
        for key, value in batch.items():
            if isinstance(value, Tensor):
                result[key] = value.to(self.device)
            elif key != "sentences":  # Keep sentences as-is
                result[key] = value
        return result

    @staticmethod
    def _result_to_full_metrics(result: Optional[EvaluationResult]) -> dict[str, Any]:
        """Convert EvaluationResult to a dict with both f1 and aligned_accuracy."""
        if result is None:
            return {}
        return {
            "f1": dict(result.f1),
            "aligned_accuracy": {k: v for k, v in result.aligned_accuracy.items() if v is not None},
            "total_words": result.total_words,
            "total_sentences": result.total_sentences,
        }

    @staticmethod
    def _log_result(metrics: dict, prefix: str, result: EvaluationResult) -> None:
        """Add F1 entries from an EvaluationResult to a metrics dict."""
        for name, value in result.f1.items():
            metrics[f"{prefix}/{name}"] = value

    def _log_epoch_metrics(
        self,
        epoch: int,
        train_loss: float,
        dev_result: EvaluationResult,
        train_result: Optional[EvaluationResult] = None,
        test_result: Optional[EvaluationResult] = None,
        dev_fulltext_result: Optional[EvaluationResult] = None,
        test_fulltext_result: Optional[EvaluationResult] = None,
    ) -> None:
        """Log epoch metrics to wandb."""
        metrics: dict[str, Any] = {
            "epoch": epoch + 1,
            "train/epoch_loss": train_loss,
        }

        self._log_result(metrics, "dev_eval", dev_result)

        if train_result:
            self._log_result(metrics, "train_eval", train_result)

        if test_result:
            self._log_result(metrics, "test_eval", test_result)

        if dev_fulltext_result:
            self._log_result(metrics, "dev_fulltext", dev_fulltext_result)

        if test_fulltext_result:
            self._log_result(metrics, "test_fulltext", test_fulltext_result)

        self.logger.log(metrics)

    def _print_epoch_summary(
        self,
        epoch: int,
        train_loss: float,
        dev_result: EvaluationResult,
        train_result: Optional[EvaluationResult],
        test_result: Optional[EvaluationResult],
        dev_fulltext_result: Optional[EvaluationResult],
        test_fulltext_result: Optional[EvaluationResult],
    ) -> None:
        """Print epoch summary to console."""
        print(f"\nEpoch {epoch + 1}/{self.config.training.epochs}")
        print(f"  Train Loss: {train_loss:.4f}")

        def _print_split(label: str, result: EvaluationResult, fulltext: Optional[EvaluationResult]) -> None:
            f1 = result.f1
            print(f"  {label} (aligned): UAS={f1.get('UAS', 0):.2f}%, LAS={f1.get('LAS', 0):.2f}%, "
                  f"UPOS={f1.get('UPOS', 0):.2f}%, Lemmas={f1.get('Lemmas', 0):.2f}%, "
                  f"MLAS={f1.get('MLAS', 0):.2f}%")
            if fulltext:
                ft = fulltext.f1
                print(f"  {label} (full-text): UAS={ft.get('UAS', 0):.2f}%, LAS={ft.get('LAS', 0):.2f}%, "
                      f"Tokens={ft.get('Tokens', 0):.2f}%, Words={ft.get('Words', 0):.2f}%, "
                      f"Sentences={ft.get('Sentences', 0):.2f}%")

        _print_split("Dev", dev_result, dev_fulltext_result)
        if test_result:
            _print_split("Test", test_result, test_fulltext_result)
        if train_result:
            _print_split("Train (subset)", train_result, None)

    def _save_checkpoint(
        self,
        epoch: int,
        metrics: dict[str, float],
        epoch_results: dict[str, Any],
    ) -> None:
        """Save model checkpoint."""
        checkpoint_dir = self.output_dir / "checkpoints"

        # Save only trainable parameters (LoRA adapters + task heads)
        trainable_names = {
            name for name, param in self.model.named_parameters()
            if param.requires_grad
        }
        trainable_state = {
            name: param for name, param in self.model.state_dict().items()
            if name in trainable_names
        }

        checkpoint = {
            "epoch": epoch,
            "model_state_dict": trainable_state,
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "metrics": metrics,
            "global_step": self.global_step,
            "best_metric": self.best_metric,
        }

        # Save epoch checkpoint
        if self.config.training.save_every_epoch:
            path = checkpoint_dir / f"checkpoint_epoch_{epoch + 1}.pt"
            torch.save(checkpoint, path)

        # Save best checkpoint
        metric_value = metrics[self.config.training.metric_for_best]
        if metric_value > self.best_metric:
            self.best_metric = metric_value
            path = checkpoint_dir / "checkpoint_best.pt"
            torch.save(checkpoint, path)
            print(f"  New best model! {self.config.training.metric_for_best}={metric_value:.2f}%")

            # Save best epoch results to a stable file
            best_results = {
                **epoch_results,
                "best_model_split": self.config.training.best_model_split,
                "metric_for_best": self.config.training.metric_for_best,
                "metric_value": metric_value,
            }
            results_path = self.output_dir / "results" / "results_best.json"
            with open(results_path, "w") as f:
                json.dump(best_results, f, indent=2, default=str)

        # Always save latest
        path = checkpoint_dir / "checkpoint_latest.pt"
        torch.save(checkpoint, path)

    def _save_results(
        self, results: dict[str, Any], prefix: str, suffix: Optional[int] = None
    ) -> None:
        """Save results to JSON file."""
        results_dir = self.output_dir / "results"
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        filename = f"results_{prefix}"
        if suffix is not None:
            filename += f"_{suffix}"
        filename += f"_{timestamp}.json"

        path = results_dir / filename
        with open(path, "w") as f:
            json.dump(results, f, indent=2, default=str)

    def _cleanup_old_checkpoints(self, checkpoint_dir: Path, keep_last: int = 5) -> None:
        """Remove old epoch checkpoints, keeping only the most recent ones."""
        epoch_checkpoints = sorted(
            checkpoint_dir.glob("checkpoint_epoch_*.pt"),
            key=lambda p: p.stat().st_mtime,
        )
        for old_ckpt in epoch_checkpoints[:-keep_last]:
            old_ckpt.unlink()

    def load_checkpoint(self, path: Path) -> None:
        """Load checkpoint to resume training.

        Args:
            path: Path to checkpoint file.
        """
        checkpoint = torch.load(path, map_location=self.device, weights_only=False)

        self.model.load_state_dict(checkpoint["model_state_dict"], strict=False)
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        self.current_epoch = checkpoint["epoch"] + 1
        self.global_step = checkpoint["global_step"]
        self.best_metric = checkpoint["best_metric"]

        print(f"Loaded checkpoint from epoch {checkpoint['epoch'] + 1}")
        print(f"  Best metric: {self.best_metric:.2f}%")
