import os
import copy
import torch
from torch.nn.utils import clip_grad_norm_

try:
    from torch.amp import autocast as _autocast, GradScaler as _GradScaler  # type: ignore
    _USE_TORCH_AMP = True
except Exception:
    from torch.cuda.amp import autocast as _autocast, GradScaler as _GradScaler  # type: ignore
    _USE_TORCH_AMP = False

from torchlitex.checkpoints import load_checkpoint, save_checkpoint
from torchlitex.dataset import create_loader
from torchlitex.ddp import setup_model
from torchlitex.logging import rank_print, init_wandb, wandb_log
from torchlitex.utils import average_tensor, move_to_device, iter_microbatches, EMA


class Trainer:
    def __init__(
        self,
        model,
        dataset,
        loss_fn,
        optimizer,
        scaler=None,
        scheduler=None,
        grad_clip_norm=None,
        amp_enabled=None,
        callbacks=None,
        eval_fn=None,
        val_dataset=None,
        log_every=None,
        wandb_project=None,
        wandb_config=None,
        wandb_run=None,
        wandb_rank: int = 0,
        gradient_accum_steps: int = 1,
        microbatch_size: int | None = None,
        ema_decay: float | None = None,
        eval_use_ema: bool = False,
    ):
        self.model = model
        self.dataset = dataset
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.grad_clip_norm = grad_clip_norm
        self.amp_enabled = torch.cuda.is_available() if amp_enabled is None else amp_enabled
        self.callbacks = callbacks or []
        self.eval_fn = eval_fn
        self.val_dataset = val_dataset
        self.log_every = log_every
        self.scaler = scaler or _GradScaler(enabled=self.amp_enabled)
        self.wandb_project = wandb_project
        self.wandb_config = wandb_config or {}
        self._wandb_run = wandb_run
        self.wandb_rank = wandb_rank
        self.gradient_accum_steps = max(1, gradient_accum_steps)
        self.microbatch_size = microbatch_size
        self.ema = EMA(self.model, decay=ema_decay) if ema_decay else None
        self.eval_use_ema = eval_use_ema
        self.global_step = 0

    def _callbacks(self, event: str, **kwargs):
        for cb in self.callbacks:
            if callable(cb):
                cb(event=event, **kwargs)
            elif hasattr(cb, event):
                getattr(cb, event)(**kwargs)

    def _forward_and_loss(self, batch):
        if isinstance(batch, dict):
            target_keys = [k for k in ("target", "targets", "label", "labels") if k in batch]
            targets = batch[target_keys[0]] if target_keys else None
            inputs = {k: v for k, v in batch.items() if k not in target_keys}
        elif isinstance(batch, (tuple, list)) and len(batch) == 2:
            inputs, targets = batch
        else:
            inputs, targets = batch, None

        if isinstance(inputs, (tuple, list)):
            outputs = self.model(*inputs)
        elif isinstance(inputs, dict):
            outputs = self.model(**inputs)
        else:
            outputs = self.model(inputs)

        if targets is None:
            if isinstance(outputs, tuple) and len(outputs) == 2:
                logits, targets = outputs
            else:
                raise ValueError(
                    "Trainer expects dataset to return (inputs, targets) or model to return (logits, targets)."
                )
        else:
            logits = outputs

        return self.loss_fn(logits, targets), logits, inputs, targets

    def ddp_train_loop(self, rank, world_size, batch_size, epochs, ckpt_path, num_workers: int = 2, drop_last: bool = False):
        train_loader = create_loader(
            self.dataset,
            batch_size=batch_size,
            rank=rank,
            world_size=world_size,
            num_workers=num_workers,
            drop_last=drop_last,
        )

        has_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
        device = torch.device(f"cuda:{rank}" if has_cuda else "cpu")

        self.model = setup_model(self.model, device)

        if rank == self.wandb_rank and (self._wandb_run or self.wandb_project):
            self._wandb_run = self._wandb_run or init_wandb(project=self.wandb_project, config=self.wandb_config)

        self._callbacks("on_train_start", trainer=self, rank=rank)

        best_val = float("inf")
        start_epoch = 0

        if os.path.exists(ckpt_path):
            start_epoch, best_val, ema_state = load_checkpoint(
                ckpt_path,
                self.model.module,
                optimizer=self.optimizer,
                scaler=self.scaler,
            )
            if self.ema and ema_state:
                self.ema.load_state_dict(ema_state)

        for epoch in range(start_epoch, epochs):
            if hasattr(train_loader, "sampler") and hasattr(train_loader.sampler, "set_epoch"):
                train_loader.sampler.set_epoch(epoch)

            self.model.train()
            total_loss = torch.tensor(0.0, device=device)
            global_step = epoch * len(train_loader)
            self.global_step = global_step
            self._callbacks("on_epoch_start", epoch=epoch, num_batches=len(train_loader), trainer=self, rank=rank)

            num_batches = len(train_loader)
            for batch_idx, batch in enumerate(train_loader):
                self._callbacks("on_batch_start", epoch=epoch, batch_idx=batch_idx, num_batches=num_batches, trainer=self, rank=rank)

                batch = move_to_device(batch, device)
                micro_iter = iter_microbatches(batch, self.microbatch_size)
                accum_counter = 0

                self.optimizer.zero_grad(set_to_none=True)

                last_logits, last_inputs, last_targets = None, None, None
                for micro in micro_iter:
                    if _USE_TORCH_AMP:
                        autocast_ctx = _autocast("cuda", enabled=self.amp_enabled and device.type == "cuda")
                    else:
                        autocast_ctx = _autocast(enabled=self.amp_enabled and device.type == "cuda")

                    with autocast_ctx:
                        loss, logits, inputs, targets = self._forward_and_loss(micro)

                    last_logits, last_inputs, last_targets = logits, inputs, targets
                    loss_scaled = loss / self.gradient_accum_steps
                    self.scaler.scale(loss_scaled).backward()
                    accum_counter += 1
                    total_loss += loss.detach()

                    should_step = accum_counter % self.gradient_accum_steps == 0

                    if should_step:
                        if self.grad_clip_norm is not None:
                            self.scaler.unscale_(self.optimizer)
                            clip_grad_norm_(self.model.parameters(), self.grad_clip_norm)

                        self.scaler.step(self.optimizer)
                        self.scaler.update()
                        self.optimizer.zero_grad(set_to_none=True)
                        global_step += 1

                        if self.ema:
                            self.ema.update(self.model.module if hasattr(self.model, "module") else self.model)

                        # allow non-tensor logits (e.g., dict) by conditional detach
                        logits_for_cb = logits.detach() if hasattr(logits, "detach") else logits
                        self._callbacks(
                            "on_step_end",
                            epoch=epoch,
                            step=global_step,
                            batch_idx=batch_idx,
                            num_batches=num_batches,
                            trainer=self,
                            rank=rank,
                            loss=float(loss.detach()),
                            logits=logits_for_cb,
                            inputs=inputs,
                            targets=targets,
                        )

                        self.global_step = global_step

                        if self.log_every and rank == 0 and global_step % self.log_every == 0:
                            rank_print(f"[Epoch {epoch} Step {global_step}] loss={float(loss):.4f}", rank=rank)
                            wandb_log({"train/loss_step": float(loss)}, step=global_step, rank=rank, active_rank=self.wandb_rank)

                # handle remainder grads
                if accum_counter % self.gradient_accum_steps != 0:
                    if self.grad_clip_norm is not None:
                        self.scaler.unscale_(self.optimizer)
                        clip_grad_norm_(self.model.parameters(), self.grad_clip_norm)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                    self.optimizer.zero_grad(set_to_none=True)
                    global_step += 1
                    if self.ema:
                        self.ema.update(self.model.module if hasattr(self.model, "module") else self.model)
                    logits_for_cb = last_logits.detach() if (last_logits is not None and hasattr(last_logits, "detach")) else last_logits
                    self._callbacks(
                        "on_step_end",
                        epoch=epoch,
                        step=global_step,
                        batch_idx=batch_idx,
                        num_batches=num_batches,
                        trainer=self,
                        rank=rank,
                        loss=float(loss.detach()),
                        logits=logits_for_cb,
                        inputs=last_inputs,
                        targets=last_targets,
                    )
                    self.global_step = global_step

                self._callbacks(
                    "on_batch_end",
                    epoch=epoch,
                    batch_idx=batch_idx,
                    num_batches=num_batches,
                    trainer=self,
                    rank=rank,
                    loss=float(loss.detach()),
                )

            total_loss = average_tensor(total_loss, device) / world_size
            best_val = min(best_val, float(total_loss))

            if self.scheduler:
                self.scheduler.step()

            # Built-in validation loop
            val_loss = None
            if self.val_dataset is not None:
                val_loader = create_loader(
                    self.val_dataset,
                    batch_size=batch_size,
                    rank=rank,
                    world_size=world_size,
                    num_workers=num_workers,
                    drop_last=False,
                    shuffle=False,
                )
                self.model.eval()
                val_total = torch.tensor(0.0, device=device)
                val_count = 0
                sample_logits, sample_inputs, sample_targets = None, None, None

                with torch.no_grad():
                    for val_batch_idx, val_batch in enumerate(val_loader):
                        val_batch = move_to_device(val_batch, device)

                        if _USE_TORCH_AMP:
                            autocast_ctx = _autocast("cuda", enabled=self.amp_enabled and device.type == "cuda")
                        else:
                            autocast_ctx = _autocast(enabled=self.amp_enabled and device.type == "cuda")

                        with autocast_ctx:
                            loss, logits, inputs, targets = self._forward_and_loss(val_batch)

                        val_total += loss.detach()
                        val_count += 1

                        # Capture first batch for sample reconstruction callback
                        if val_batch_idx == 0:
                            if isinstance(logits, dict):
                                sample_logits = {k: v.detach() if hasattr(v, 'detach') else v for k, v in logits.items()}
                            else:
                                sample_logits = logits.detach() if hasattr(logits, 'detach') else logits
                            sample_inputs, sample_targets = inputs, targets

                val_loss = average_tensor(val_total, device) / (val_count * world_size) if val_count > 0 else torch.tensor(0.0, device=device)
                val_loss = float(val_loss)

                if rank == 0:
                    rank_print(f"[Epoch {epoch}] val_loss={val_loss:.4f}", rank=rank)
                    wandb_log({"val/loss_epoch": val_loss}, step=epoch, rank=rank, active_rank=self.wandb_rank)

                self._callbacks(
                    "on_val_end",
                    epoch=epoch,
                    val_loss=val_loss,
                    trainer=self,
                    rank=rank,
                    sample_logits=sample_logits,
                    sample_inputs=sample_inputs,
                    sample_targets=sample_targets,
                )

                self.model.train()

            if rank == 0 and self.eval_fn:
                with torch.no_grad():
                    eval_model = self.model.module if hasattr(self.model, "module") else self.model
                    if self.eval_use_ema and self.ema:
                        backup = {k: v.detach().clone() for k, v in eval_model.state_dict().items()}
                        self.ema.apply_to(eval_model)
                    eval_score = self.eval_fn(
                        eval_model,
                        device=device,
                        epoch=epoch,
                        rank=rank,
                        world_size=world_size,
                    )
                    if self.eval_use_ema and self.ema:
                        eval_model.load_state_dict(backup, strict=False)
                best_val = min(best_val, float(eval_score))

            if rank == 0:
                save_checkpoint(
                    ckpt_path,
                    self.model,
                    self.optimizer,
                    self.scaler,
                    epoch,
                    best_val,
                    ema_state=self.ema.state_dict() if self.ema else None,
                )
                rank_print(f"[Epoch {epoch}] loss={float(total_loss):.4f}", rank=rank)
                wandb_log(
                    {"train/loss_epoch": float(total_loss), "train/best": best_val, "train/epoch": epoch},
                    step=epoch,
                    rank=rank,
                    active_rank=self.wandb_rank,
                )

                self._callbacks("on_epoch_end", epoch=epoch, loss=float(total_loss), val_loss=val_loss, best=best_val, trainer=self, rank=rank)

        if rank == 0:
            rank_print("Training complete.", rank=rank)
            self._callbacks("on_train_end", trainer=self, rank=rank)


class TQDMCallback:
    """
    Optional tqdm progress bar callback.
    """

    def __init__(self):
        self.pbar = None

    def on_epoch_start(self, epoch, num_batches, rank, **_):
        if rank != 0:
            return
        try:
            from tqdm import tqdm  # type: ignore
        except Exception:
            return
        self.pbar = tqdm(total=num_batches, desc=f"Epoch {epoch}")

    def on_batch_end(self, rank, **_):
        if rank != 0 or self.pbar is None:
            return
        self.pbar.update(1)

    def on_epoch_end(self, rank, **_):
        if rank != 0 or self.pbar is None:
            return
        self.pbar.close()
        self.pbar = None
