# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates
# Copyright 2026 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""VLM Trainer assembled from the shared BaseTrainer stages."""

from collections import defaultdict
from typing import Any, Dict

from hyper_parallel import SkipDTensorDispatch
from hyper_parallel.core.utils import clip_grad_norm_
from hyper_parallel.data.batching import calculate_num_micro_batches
from hyper_parallel.data.vlm import build_processor, build_vlm_get_batch
from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
from hyper_parallel.trainer.runtime.logging import create_logger
from hyper_parallel.trainer.runtime.memory import print_device_mem_info
from hyper_parallel.trainer.runtime.device import synchronize  # pylint: disable=syntax-error
from hyper_parallel.trainer.base import BaseTrainer
from hyper_parallel.trainer.config import TrainerConfig

logger = create_logger(__name__)


class VLMTrainer:
    """Compose the VLM training runtime from explicit BaseTrainer stages."""

    base: BaseTrainer

    def __init__(self, config: TrainerConfig) -> None:
        """Build the VLM Trainer in the same explicit order as VeOmni."""
        self.base = BaseTrainer.__new__(BaseTrainer)
        self.base.config = config

        self.base._setup()
        self.base._build_model()
        self.base._build_loss()

        # datasets
        self._build_model_assets()
        self._build_data_transform()
        self.base._build_dataset()

        # dataloader
        self._build_collate_fn()
        self.base._build_dataloader()

        # get_batch
        self._build_get_batch()
        self.base._compute_train_iters()

        self.base._build_optimizer()
        self.base._build_lr_scheduler()
        self.base._build_training_context()
        self.base._init_callbacks()

    def _build_model_assets(self) -> None:
        """Build processor-backed assets for VLM training."""
        config: TrainerConfig = self.base.config
        if config.dataset is None:
            raise ValueError("dataset must define a build target")

        processor_path = (
            getattr(config.model, "tokenizer_path", None)
            or getattr(config.model, "pretrained_model_name_or_path", None)
        )
        if processor_path is None:
            self.base.processor = None
        else:
            self.base.processor = build_processor(processor_path)
        self.base.tokenizer = getattr(self.base.processor, "tokenizer", None)
        self.base.chat_template = None

        self.base.model_assets = [self.base.model_config]
        if self.base.processor is not None:
            self.base.model_assets.append(self.base.processor)

    def _build_data_transform(self) -> None:
        """Build the configured multimodal sample transform."""
        dataset_config = self.base.config.dataset
        if dataset_config is None:
            raise ValueError("dataset must define a build target")
        if dataset_config.data_transform is None:
            self.base.data_transform = None
            return
        self.base.data_transform = dataset_config.data_transform.build(
            processor=self.base.processor,
        )

    def _build_collate_fn(self) -> None:
        """Build the VLM collator and gradient-accumulation batch count."""
        dataloader_config = self.base.config.dataloader
        if dataloader_config is None or dataloader_config.collate_fn is None:
            raise ValueError("dataloader.collate_fn must define a build target")
        training_config = self.base.config.training
        self.base.num_micro_batches = calculate_num_micro_batches(
            global_batch_size=training_config.global_batch_size,
            micro_batch_size=training_config.micro_batch_size,
            dp_world_size=self.base.mesh.dp_size,
        )
        self.base.collate_fn = dataloader_config.collate_fn.build()

    def _build_get_batch(self) -> None:
        """Build the DataLoader-to-VLM batch adapter."""
        config = self.base.config
        get_batch_builder = (
            config.dataloader.get_batch.build
            if config.dataloader.get_batch
            else build_vlm_get_batch
        )
        self.base.get_batch = get_batch_builder(
            mesh_context=self.base.mesh,
            device=self.base.device,
            pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
        )

    @property
    def distributed_setup(self) -> Any:
        """Return the shared distributed setup."""
        return self.base.distributed_setup

    @property
    def mesh(self) -> Any:
        """Return the shared mesh context."""
        return self.base.mesh

    @property
    def dp_cp_mesh(self) -> Any:
        """Return the shared data/context-parallel mesh."""
        return self.base.dp_cp_mesh

    def on_train_begin(self) -> None:
        """Dispatch the training-begin lifecycle hook."""
        self.base.on_train_begin()

    def on_train_end(self) -> None:
        """Dispatch the training-end lifecycle hook."""
        self.base.on_train_end()

    def on_epoch_begin(self) -> None:
        """Dispatch the epoch-begin lifecycle hook."""
        self.base.on_epoch_begin()

    def on_epoch_end(self) -> None:
        """Dispatch the epoch-end lifecycle hook."""
        self.base.on_epoch_end()

    def on_step_begin(self, micro_batches: Any = None) -> None:
        """Dispatch the step-begin lifecycle hook."""
        self.base.on_step_begin(micro_batches=micro_batches)

    def on_step_end(
            self,
            loss: Any = None,
            loss_dict: Any = None,
            grad_norm: Any = None,
    ) -> None:
        """Dispatch the step-end lifecycle hook."""
        self.base.on_step_end(
            loss=loss,
            loss_dict=loss_dict,
            grad_norm=grad_norm,
        )

    def train_step(self, data_iterator: Any) -> Dict[str, float]:
        """Execute one VLM training step."""
        config = self.base.config
        first_training_batch = self.base.get_batch(data_iterator)
        num_micro_steps = self.base.num_micro_batches
        training_batches = [first_training_batch]
        for _ in range(1, num_micro_steps):
            training_batches.append(self.base.get_batch(data_iterator))

        self.on_step_begin(
            micro_batches=[model_inputs for model_inputs, _ in training_batches]
        )
        synchronize()

        total_loss = 0.0
        total_loss_dict = defaultdict(int)

        for micro_step, batch in enumerate(training_batches):
            model_inputs, loss_inputs = batch
            self.base.model_reshard(micro_step, num_micro_steps)
            self.base.configure_fsdp_gradient_sync(
                micro_step,
                num_micro_steps,
            )
            self.base.current_token_counts = count_loss_token(loss_inputs)
            self.base.step_token_counts = {
                name: token_count * num_micro_steps
                for name, token_count in self.base.current_token_counts.items()
            }
            loss, loss_dict = self.base.forward_backward_step(model_inputs)

            # Release each device batch as soon as its backward pass completes;
            # prefetching all micro-batches must not pin them until step end.
            training_batches[micro_step] = None

            total_loss += loss.item()
            for loss_name, loss_value in loss_dict.items():
                total_loss_dict[loss_name] += loss_value.item()

        grad_norm = clip_grad_norm_(
            self.base.model,
            config.training.max_grad_norm,
        )

        optimizers = (
            self.base.optimizer
            if isinstance(self.base.optimizer, list)
            else [self.base.optimizer]
        )
        for optimizer in optimizers:
            with SkipDTensorDispatch():
                optimizer.step()
            optimizer.zero_grad()

        schedulers = (
            self.base.lr_scheduler
            if isinstance(self.base.lr_scheduler, list)
            else (
                [self.base.lr_scheduler]
                if self.base.lr_scheduler is not None
                else []
            )
        )
        for scheduler in schedulers:
            scheduler.step()

        grad_norm_value = float(grad_norm)
        self.on_step_end(
            loss=total_loss,
            loss_dict=total_loss_dict,
            grad_norm=grad_norm_value,
        )

        self.base.state.global_step += 1

        return {
            "loss": total_loss,
            "grad_norm": grad_norm_value,
        }

    def train(self) -> None:
        """Run the VLM training loop."""
        config = self.base.config
        self.on_train_begin()
        logger.info(
            "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
            self.base.local_rank,
            self.base.state.global_step,
            self.base.train_iters,
            self.base.state.epoch,
            self.base.train_epochs,
        )

        # Checkpoint resume restores state.global_step, state.epoch, and the DataLoader cursor.
        for epoch in range(self.base.state.epoch, self.base.train_epochs):
            train_dataloader = self.base.train_dataloader

            if hasattr(train_dataloader, "set_epoch"):
                train_dataloader.set_epoch(epoch)

            self.on_epoch_begin()
            data_iterator = iter(train_dataloader) if train_dataloader is not None else None

            start_step = self.base.state.global_step - epoch * self.base.train_steps
            train_steps = min(self.base.train_steps, self.base.train_iters - epoch * self.base.train_steps)
            for _ in range(start_step, train_steps):
                try:
                    self.train_step(data_iterator)
                except StopIteration:
                    logger.info("epoch:%s Dataloader finished with drop_last %s", epoch, config.dataloader.drop_last)
                    break

            self.on_epoch_end()
            self.base.state.epoch = epoch + 1
            print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")

            if self.base.state.global_step >= self.base.train_iters:
                break

        self.on_train_end()

        synchronize()
        self.base.destroy_distributed()


__all__ = ["VLMTrainer"]