from __future__ import annotations

from abc import ABC, abstractmethod
from pathlib import Path
import pprint
import time
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.amp import GradScaler

from lib.data import build_dataloader
from lib.tools import build_metrics, compute_metrics, reset_metrics, update_metrics

try:
    import swanlab
except ImportError:
    swanlab = None


class BaseTrainer(ABC):
    """
    训练器基类。

    设计目标：
    - 统一配置入口
    - 统一模型/优化器/调度器创建
    - 不同训练流程只重写最少的方法
    """

    def __init__(self, cfg: dict[str, Any], args: Any | None = None) -> None:
        self.cfg = cfg
        self.args = args
        self.device = self._build_device()
        self.output_dir = self._build_output_dir()
        self.start_epoch = 0
        self.best_metric: float | None = None
        self.no_improve_epochs = 0
        self.swanlab_run = None
        self.grad_scaler = GradScaler("cuda", enabled=self._amp_enabled())

    def _build_device(self) -> torch.device:
        device_name = self.cfg.get("train", {}).get("device", "cpu")
        if device_name == "cuda" and not torch.cuda.is_available():
            device_name = "cpu"
        return torch.device(device_name)

    def _build_output_dir(self) -> Path:
        output_dir = self.cfg.get("checkpoint", {}).get("dir", "outputs/supervised_segmentation")
        path = Path(output_dir)
        path.mkdir(parents=True, exist_ok=True)
        return path

    def _amp_enabled(self) -> bool:
        return bool(self.cfg.get("train", {}).get("amp", False)) and self.device.type == "cuda"

    def _auto_batch_size_cfg(self) -> dict[str, Any]:
        cfg = self.cfg.get("train", {}).get("auto_batch_size", {})
        return cfg if isinstance(cfg, dict) else {}

    def _auto_batch_size_enabled(self) -> bool:
        return bool(self._auto_batch_size_cfg().get("enabled", False))

    def _gpu_total_memory_gb(self) -> float | None:
        if self.device.type != "cuda" or not torch.cuda.is_available():
            return None
        props = torch.cuda.get_device_properties(self.device)
        return float(props.total_memory / (1024 ** 3))

    def _estimate_auto_batch_size(self, *, default_batch_size: int, ssl: bool = False) -> int:
        cfg = self._auto_batch_size_cfg()
        if not cfg.get("enabled", False):
            return int(default_batch_size)

        total_gb = self._gpu_total_memory_gb()
        if total_gb is None:
            return int(default_batch_size)

        target_fraction = float(cfg.get("target_memory_fraction", 0.75))
        target_fraction = min(max(target_fraction, 0.1), 0.95)
        reference_gpu_gb = float(cfg.get("reference_gpu_gb", 8.0))
        reference_batch_size = int(cfg.get("reference_batch_size", default_batch_size))
        max_batch_size = int(cfg.get("max_batch_size", reference_batch_size))
        min_batch_size = int(cfg.get("min_batch_size", 1))

        memory_penalty = float(cfg.get("memory_penalty", 1.0 if not ssl else 1.35))
        scaled = int((reference_batch_size * total_gb * target_fraction) / max(reference_gpu_gb * 0.75 * memory_penalty, 1e-6))
        batch_size = max(min_batch_size, min(max_batch_size, max(default_batch_size, scaled)))
        return int(batch_size)

    def _resolve_batch_size(self, key: str, default: int, *, ssl: bool = False) -> int:
        train_cfg = self.cfg.get("train", {})
        configured = int(train_cfg.get(key, default))
        batch_size = self._estimate_auto_batch_size(default_batch_size=configured, ssl=ssl)
        if self._auto_batch_size_enabled() and batch_size != configured:
            print(
                {
                    "message": "auto_batch_size adjusted",
                    "key": key,
                    "configured": configured,
                    "resolved": batch_size,
                    "gpu_total_gb": self._gpu_total_memory_gb(),
                }
            )
        return batch_size

    def _dataset_cfg(self) -> dict[str, Any]:
        return self.cfg.get("dataset", {})

    def _dataset_name(self) -> str:
        dataset_cfg = self._dataset_cfg()
        dataset_name = dataset_cfg.get("dataset_name") or dataset_cfg.get("name")
        if not dataset_name:
            raise ValueError("dataset.dataset_name is required.")
        return str(dataset_name)

    def _dataset_root(self) -> str:
        dataset_cfg = self._dataset_cfg()
        root = dataset_cfg.get("root")
        if not root:
            raise ValueError("dataset.root is required.")
        return str(root)

    def _image_size(self) -> tuple[int, int]:
        dataset_cfg = self._dataset_cfg()
        image_size = dataset_cfg.get("image_size")
        if image_size is None:
            raise ValueError("dataset.image_size is required.")
        return int(image_size[0]), int(image_size[1])

    def _build_resize_transform(self, *, mode: str) -> Any:
        height, width = self._image_size()
        interpolation_mode = "bilinear" if mode == "image" else "nearest"

        def _transform(tensor: torch.Tensor) -> torch.Tensor:
            resized = F.interpolate(
                tensor.unsqueeze(0),
                size=(height, width),
                mode=interpolation_mode,
                align_corners=False if interpolation_mode != "nearest" else None,
            )
            return resized.squeeze(0)

        return _transform

    def _build_segmentation_loader(
            self,
            *,
            split: str,
            batch_size: int,
            shuffle: bool,
            split_file: str | None = None,
    ):
        dataset_cfg = self._dataset_cfg()
        train_cfg = self.cfg.get("train", {})
        num_workers = max(0, int(train_cfg.get("num_workers", 0)))
        persistent_workers = bool(train_cfg.get("persistent_workers", False)) if num_workers > 0 else False
        loader = build_dataloader(
            dataset_name=self._dataset_name(),
            root=self._dataset_root(),
            split=split,
            split_file=split_file,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            image_transform=self._build_resize_transform(mode="image"),
            mask_transform=self._build_resize_transform(mode="mask"),
            pin_memory=bool(train_cfg.get("pin_memory", self.device.type == "cuda")),
            persistent_workers=persistent_workers,
            prefetch_factor=train_cfg.get("prefetch_factor") if num_workers > 0 else None,
        )
        return loader

    def _build_val_loader(
            self,
            *,
            batch_size: int,
            shuffle: bool = False,
    ):
        dataset_cfg = self._dataset_cfg()
        val_split = dataset_cfg.get("val_split", "val")
        if val_split is None:
            return None
        return self._build_segmentation_loader(
            split=str(val_split),
            split_file=dataset_cfg.get("val_split_file"),
            batch_size=batch_size,
            shuffle=shuffle,
        )

    def _checkpoint_cfg(self) -> dict[str, Any]:
        return self.cfg.get("checkpoint", {})

    def _logging_cfg(self) -> dict[str, Any]:
        return self.cfg.get("logging", {})

    def _validation_cfg(self) -> dict[str, Any]:
        return self.cfg.get("validation", {})

    def _checkpoint_enabled(self) -> bool:
        return bool(self._checkpoint_cfg().get("save", True))

    def _best_mode(self) -> str:
        return str(self._checkpoint_cfg().get("monitor_mode", "min"))

    def _is_better_metric(self, metric: float) -> bool:
        if self.best_metric is None:
            return True
        if self._best_mode() == "max":
            return metric > self.best_metric
        return metric < self.best_metric

    def _save_checkpoint(self, filename: str, state: dict[str, Any]) -> Path | None:
        if not self._checkpoint_enabled():
            return None
        path = self.output_dir / filename
        torch.save(state, path)
        return path

    def _resume_checkpoint_path(self) -> Path | None:
        resume_path = self._checkpoint_cfg().get("resume")
        if not resume_path:
            return None
        path = Path(str(resume_path))
        if not path.is_absolute():
            path = Path.cwd() / path
        return path

    def _maybe_resume(
            self,
            *,
            module_map: dict[str, Any],
            optimizer: Any | None = None,
            scheduler: Any | None = None,
    ) -> dict[str, Any] | None:
        path = self._resume_checkpoint_path()
        if path is None:
            return None
        if not path.exists():
            raise FileNotFoundError(f"Resume checkpoint not found: {path}")

        checkpoint = torch.load(path, map_location="cpu")
        strict = bool(self._checkpoint_cfg().get("resume_strict", True))
        for key, module in module_map.items():
            if module is None:
                continue
            state_dict = checkpoint.get(key)
            if state_dict is not None:
                module.load_state_dict(state_dict, strict=strict)

        if optimizer is not None and checkpoint.get("optimizer") is not None:
            optimizer.load_state_dict(checkpoint["optimizer"])
        if scheduler is not None and checkpoint.get("scheduler") is not None:
            scheduler.load_state_dict(checkpoint["scheduler"])
        if checkpoint.get("grad_scaler") is not None:
            self.grad_scaler.load_state_dict(checkpoint["grad_scaler"])

        if checkpoint.get("best_metric") is not None:
            self.best_metric = float(checkpoint["best_metric"])
        elif checkpoint.get("metrics") is not None:
            monitor_name = str(self._checkpoint_cfg().get("monitor", "total"))
            monitor_value = checkpoint["metrics"].get(f"val_{monitor_name}")
            if monitor_value is None:
                monitor_value = checkpoint["metrics"].get(monitor_name)
            if monitor_value is not None:
                self.best_metric = float(monitor_value)
        if checkpoint.get("no_improve_epochs") is not None:
            self.no_improve_epochs = int(checkpoint["no_improve_epochs"])

        if bool(self._checkpoint_cfg().get("resume_training", True)):
            self.start_epoch = int(checkpoint.get("epoch", -1)) + 1
        return checkpoint

    def _validation_enabled(self) -> bool:
        return bool(self._validation_cfg().get("enabled", True))

    def _validation_interval(self) -> int:
        return max(1, int(self._validation_cfg().get("interval", 1)))

    def _should_validate(self, epoch: int) -> bool:
        return self._validation_enabled() and ((epoch + 1) % self._validation_interval() == 0)

    def _metric_task_mode(self) -> str:
        validation_cfg = self._validation_cfg()
        metrics_cfg = validation_cfg.get("metrics", self.cfg.get("metrics"))
        if isinstance(metrics_cfg, dict):
            return str(metrics_cfg.get("task_mode", "binary"))
        return "binary"

    def _metric_threshold(self) -> float:
        validation_cfg = self._validation_cfg()
        threshold = validation_cfg.get("threshold", 0.5)
        return float(threshold)

    def _build_validation_metrics(self) -> dict[str, Any]:
        validation_cfg = self._validation_cfg()
        metrics_cfg = validation_cfg.get("metrics", self.cfg.get("metrics"))
        if metrics_cfg is None:
            return {}
        return build_metrics(metrics_cfg)

    def _early_stopping_enabled(self) -> bool:
        return bool(self._validation_cfg().get("early_stopping", False))

    def _early_stopping_patience(self) -> int:
        return max(1, int(self._validation_cfg().get("early_stopping_patience", 10)))

    def _early_stopping_min_delta(self) -> float:
        return float(self._validation_cfg().get("early_stopping_min_delta", 0.0))

    def _update_validation_metrics(
            self,
            metrics: dict[str, Any],
            *,
            logits: torch.Tensor,
            target: torch.Tensor,
    ) -> None:
        if not metrics:
            return
        update_metrics(
            metrics,
            logits,
            target,
            task_mode=self._metric_task_mode(),
            threshold=self._metric_threshold(),
            num_classes=int(self._dataset_cfg().get("num_classes", 1)),
        )

    def _compute_validation_metric_values(self, metrics: dict[str, Any]) -> dict[str, float]:
        if not metrics:
            return {}
        values = compute_metrics(metrics)
        reset_metrics(metrics)
        return values

    def _init_swanlab(self) -> None:
        logging_cfg = self._logging_cfg()
        if not bool(logging_cfg.get("use_swanlab", False)):
            return
        if swanlab is None:
            print("SwanLab is not installed. Logging will continue without SwanLab.")
            return

        run_name = logging_cfg.get("experiment_name") or self.output_dir.name
        self.swanlab_run = swanlab.init(
            project=logging_cfg.get("project", "X_SSL_Net"),
            name=run_name,
            config=self.cfg,
            mode=logging_cfg.get("swanlab_mode"),
        )

    def _log_metrics(self, metrics: dict[str, float], *, step: int) -> None:
        if self.swanlab_run is None:
            return
        swanlab.log(metrics, step=step)

    def _close_loggers(self) -> None:
        if self.swanlab_run is not None:
            swanlab.finish()
            self.swanlab_run = None

    def _log_interval(self) -> int:
        return max(1, int(self._logging_cfg().get("log_interval", 20)))

    def _grad_clip_cfg(self) -> dict[str, Any]:
        cfg = self.cfg.get("train", {}).get("grad_clip", {})
        return cfg if isinstance(cfg, dict) else {}

    def _grad_clip_enabled(self) -> bool:
        return bool(self._grad_clip_cfg().get("enabled", False))

    def _clip_gradients(self, module: nn.Module | None) -> float | None:
        if module is None or not self._grad_clip_enabled():
            return None
        cfg = self._grad_clip_cfg()
        max_norm = float(cfg.get("max_norm", 1.0))
        norm_type = float(cfg.get("norm_type", 2.0))
        params = [param for param in module.parameters() if param.requires_grad and param.grad is not None]
        if not params:
            return None
        total_norm = torch.nn.utils.clip_grad_norm_(params, max_norm=max_norm, norm_type=norm_type)
        return float(total_norm.detach().cpu() if isinstance(total_norm, torch.Tensor) else total_norm)

    def _current_lrs(self, optimizer: Any | None) -> list[float]:
        if optimizer is None:
            return []
        return [float(group.get("lr", 0.0)) for group in optimizer.param_groups]

    @staticmethod
    def _count_parameters(module: nn.Module | None) -> dict[str, int]:
        if module is None:
            return {"total": 0, "trainable": 0}
        total = sum(param.numel() for param in module.parameters())
        trainable = sum(param.numel() for param in module.parameters() if param.requires_grad)
        return {"total": int(total), "trainable": int(trainable)}

    @staticmethod
    def _loader_summary(loader: Any | None) -> dict[str, Any] | None:
        if loader is None:
            return None
        dataset = getattr(loader, "dataset", None)
        return {
            "dataset_size": len(dataset) if dataset is not None else None,
            "num_batches": len(loader),
            "batch_size": getattr(loader, "batch_size", None),
            "num_workers": getattr(loader, "num_workers", None),
            "pin_memory": getattr(loader, "pin_memory", None),
            "persistent_workers": getattr(loader, "persistent_workers", None),
            "prefetch_factor": getattr(loader, "prefetch_factor", None),
            "drop_last": getattr(loader, "drop_last", None),
        }

    def _training_setup_summary(
            self,
            *,
            model_map: dict[str, nn.Module | None],
            loader_map: dict[str, Any | None],
            optimizer: Any | None = None,
            scheduler: Any | None = None,
    ) -> dict[str, Any]:
        return {
            "trainer": self.cfg.get("trainer", {}).get("name"),
            "device": str(self.device),
            "amp_enabled": self._amp_enabled(),
            "output_dir": str(self.output_dir),
            "start_epoch": self.start_epoch,
            "train": self.cfg.get("train", {}),
            "dataset": self.cfg.get("dataset", {}),
            "model": self.cfg.get("model", {}),
            "optimizer": self.cfg.get("optimizer", {}),
            "scheduler": self.cfg.get("scheduler"),
            "current_lrs": self._current_lrs(optimizer),
            "validation": self.cfg.get("validation", {}),
            "checkpoint": self.cfg.get("checkpoint", {}),
            "logging": self.cfg.get("logging", {}),
            "model_parameters": {
                name: self._count_parameters(module)
                for name, module in model_map.items()
            },
            "loaders": {
                name: self._loader_summary(loader)
                for name, loader in loader_map.items()
            },
            "cuda": {
                "available": torch.cuda.is_available(),
                "device_name": torch.cuda.get_device_name(self.device) if self.device.type == "cuda" else None,
                "device_count": torch.cuda.device_count(),
            },
        }

    def _print_training_setup(
            self,
            *,
            model_map: dict[str, nn.Module | None],
            loader_map: dict[str, Any | None],
            optimizer: Any | None = None,
            scheduler: Any | None = None,
    ) -> None:
        if not bool(self._logging_cfg().get("print_training_setup", True)):
            return
        summary = self._training_setup_summary(
            model_map=model_map,
            loader_map=loader_map,
            optimizer=optimizer,
            scheduler=scheduler,
        )
        print("========== TRAINING SETUP ==========")
        pprint.pprint(summary, sort_dicts=False, width=120)
        print("======== END TRAINING SETUP ========")

    def _gpu_memory_mb(self) -> float:
        if self.device.type != "cuda" or not torch.cuda.is_available():
            return 0.0
        return float(torch.cuda.max_memory_allocated(device=self.device) / (1024 ** 2))

    def _performance_snapshot(
            self,
            *,
            epoch: int,
            step: int,
            num_steps: int,
            data_time: float,
            iter_time: float,
            metrics: dict[str, float],
            prefix: str = "train",
    ) -> dict[str, float | int]:
        snapshot: dict[str, float | int] = {
            "epoch": epoch,
            "step": step,
            "num_steps": num_steps,
            "data_time": data_time,
            "iter_time": iter_time,
            "gpu_memory_mb": self._gpu_memory_mb(),
        }
        lrs = self._current_lrs(getattr(self, "optimizer", None))
        if lrs:
            snapshot["lr"] = lrs[0]
        for key, value in metrics.items():
            snapshot[f"{prefix}_{key}"] = value
        return snapshot

    def _maybe_log_step(
            self,
            *,
            epoch: int,
            step: int,
            num_steps: int,
            data_time: float,
            iter_time: float,
            metrics: dict[str, float],
            prefix: str = "train",
    ) -> None:
        if step % self._log_interval() != 0 and step != num_steps:
            return
        snapshot = self._performance_snapshot(
            epoch=epoch,
            step=step,
            num_steps=num_steps,
            data_time=data_time,
            iter_time=iter_time,
            metrics=metrics,
            prefix=prefix,
        )
        print(snapshot)
        log_metrics = {
            f"{prefix}/{key}": value
            for key, value in metrics.items()
        }
        log_metrics.update(
            {
                f"{prefix}/data_time": data_time,
                f"{prefix}/iter_time": iter_time,
                f"{prefix}/gpu_memory_mb": float(snapshot["gpu_memory_mb"]),
            }
        )
        if "lr" in snapshot:
            log_metrics[f"{prefix}/lr"] = float(snapshot["lr"])
        self._log_metrics(log_metrics, step=epoch * max(1, num_steps) + step)

    @staticmethod
    def _average_metric_sums(metric_sums: dict[str, float], steps: int) -> dict[str, float]:
        if steps <= 0:
            return {}
        return {key: value / steps for key, value in metric_sums.items()}

    def _base_checkpoint_state(self, *, epoch: int, metrics: dict[str, float] | None = None) -> dict[str, Any]:
        state = {
            "epoch": epoch,
            "cfg": self.cfg,
            "metrics": metrics or {},
            "grad_scaler": self.grad_scaler.state_dict(),
            "no_improve_epochs": self.no_improve_epochs,
        }
        return state

    def _finalize_epoch(
            self,
            *,
            epoch: int,
            train_metrics: dict[str, float],
            val_metrics: dict[str, float] | None,
            checkpoint_state: dict[str, Any],
    ) -> tuple[dict[str, Any], bool]:
        merged_metrics = dict(train_metrics)
        if val_metrics is not None:
            merged_metrics.update({f"val_{key}": value for key, value in val_metrics.items()})

        improved = False
        if val_metrics is not None:
            monitor_name = str(self._checkpoint_cfg().get("monitor", "total"))
            if monitor_name not in val_metrics:
                raise KeyError(f"Checkpoint monitor '{monitor_name}' not found in val metrics.")
            monitor_value = float(val_metrics[monitor_name])
            delta = self._early_stopping_min_delta()
            previous_best = self.best_metric
            is_better = self._is_better_metric(monitor_value)
            if previous_best is not None and self._best_mode() == "max":
                is_better = monitor_value > (previous_best + delta)
            elif previous_best is not None and self._best_mode() == "min":
                is_better = monitor_value < (previous_best - delta)

            if is_better:
                self.best_metric = monitor_value
                self.no_improve_epochs = 0
                improved = True
                best_state = dict(checkpoint_state)
                best_state.update(
                    self._base_checkpoint_state(
                        epoch=epoch,
                        metrics=merged_metrics,
                    )
                )
                best_state["best_metric"] = self.best_metric
                self._save_checkpoint("best.pth", best_state)
            else:
                self.no_improve_epochs += 1

        save_last = bool(self._checkpoint_cfg().get("save_last", True))
        if save_last:
            last_state = dict(checkpoint_state)
            last_state.update(self._base_checkpoint_state(epoch=epoch, metrics=merged_metrics))
            if self.best_metric is not None:
                last_state["best_metric"] = self.best_metric
            self._save_checkpoint("last.pth", last_state)

        summary = {"epoch": epoch}
        summary.update(train_metrics)
        if val_metrics is not None:
            summary.update({f"val_{key}": value for key, value in val_metrics.items()})
        if self.best_metric is not None:
            summary["best_metric"] = float(self.best_metric)
        summary["no_improve_epochs"] = self.no_improve_epochs
        lrs = self._current_lrs(getattr(self, "optimizer", None))
        if lrs:
            summary["lr"] = lrs[0]
        self._log_metrics(summary, step=epoch)
        should_stop = False
        if val_metrics is not None and self._early_stopping_enabled():
            should_stop = self.no_improve_epochs >= self._early_stopping_patience()
            summary["early_stop"] = should_stop
            summary["improved"] = improved
        return summary, should_stop

    @abstractmethod
    def build(self) -> None:
        """
        创建模型、优化器、数据加载器等运行所需对象。
        """

    @abstractmethod
    def train(self) -> None:
        """
        执行完整训练流程。
        """