# 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.
"""Text Trainer assembled from the shared BaseTrainer stages."""

from collections import defaultdict
from typing import Any, Dict

import torch  # pylint: disable=forbidden-backend-import

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.text import build_chat_template
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
from hyper_parallel.trainer.base import BaseTrainer
from hyper_parallel.trainer.config import TrainerConfig

logger = create_logger(__name__)


class TextTrainer:
    """Compose the text training runtime from explicit BaseTrainer stages."""

    base: BaseTrainer

    def __init__(self, config: TrainerConfig) -> None:
        """Build the text 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 tokenizer-backed assets for text training."""
        config: TrainerConfig = self.base.config
        if config.dataset is None:
            raise ValueError("dataset must define a build target")
        assets_config = config.dataset.model_assets
        tokenizer_target = assets_config.tokenizer
        if tokenizer_target is None:
            self.base.tokenizer = None
        elif getattr(tokenizer_target, "pretrained_model_name_or_path", None):
            self.base.tokenizer = tokenizer_target.build()
        else:
            self.base.tokenizer = tokenizer_target.build(
                pretrained_model_name_or_path=config.model.pretrained_model_name_or_path,
            )

        self.base.chat_template = None
        if assets_config.chat_template is not None:
            if self.base.tokenizer is None:
                raise ValueError("dataset.model_assets.tokenizer is required for conversation data")
            if isinstance(assets_config.chat_template, str):
                self.base.chat_template = build_chat_template(
                    assets_config.chat_template,
                    self.base.tokenizer,
                )
            else:
                self.base.chat_template = assets_config.chat_template.build(
                    tokenizer=self.base.tokenizer,
                )

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

    def _build_data_transform(self) -> None:
        """Build the configured text 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(
            tokenizer=self.base.tokenizer,
            chat_template=self.base.chat_template,
        )

    def _build_collate_fn(self) -> None:
        """Build the text 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(
            mesh_context=self.base.mesh,
        )

    def _build_get_batch(self) -> None:
        """Build the DataLoader-to-LLM batch adapter."""
        config = self.base.config
        if config.dataloader.get_batch is None:
            raise ValueError("dataloader.get_batch must define a batching runtime target")
        get_batch = config.dataloader.get_batch.build(
            mesh_context=self.base.mesh,
            device=self.base.device,
            tokenizer=self.base.tokenizer,
            data_config=getattr(config.dataset, "data_config", {}),
            pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
        )
        self.base.get_batch = get_batch

    @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) -> None:
        """Dispatch the step-begin lifecycle hook."""
        self.base.on_step_begin()

    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 forward_backward_step(
            self,
            data_iterator: Any,
            num_micro_steps: int,
    ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """Fetch and execute one forward-backward micro-step.

        Args:
            data_iterator: Iterator providing one raw FB batch per call.
            num_micro_steps: Number of FB steps in the current optimizer step.

        Returns:
            Loss tensor and named loss tensors for this FB step.
        """
        model_inputs, loss_inputs = self.base.get_batch(data_iterator)
        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, loss_inputs)

        return loss, loss_dict

    def train_step(self, data_iterator: Any) -> Dict[str, float]:
        """Execute one text training step."""
        config = self.base.config
        num_micro_steps = self.base.num_micro_batches
        optimizers = self.base.optimizer if isinstance(self.base.optimizer, list) else [self.base.optimizer]

        self.on_step_begin()
        synchronize()

        total_loss = 0.0
        total_loss_dict = defaultdict(int)

        for micro_step in range(num_micro_steps):
            self.base.model_reshard(micro_step, num_micro_steps)
            self.base.configure_fsdp_gradient_sync(
                micro_step,
                num_micro_steps,
            )
            loss, loss_dict = self.forward_backward_step(
                data_iterator,
                num_micro_steps,
            )

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

        grad_norm = 0.0
        if config.training.max_grad_norm > 0:
            grad_norm = clip_grad_norm_(
                self.base.model,
                config.training.max_grad_norm,
            )

        for optimizer in optimizers:
            with SkipDTensorDispatch(no_skip={torch.zeros_like}):
                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()

        # Checkpoint and logging callbacks observe the number of completed
        # optimizer updates.
        self.base.state.global_step += 1
        grad_norm_value = float(grad_norm)
        self.on_step_end(
            loss=total_loss,
            loss_dict=total_loss_dict,
            grad_norm=grad_norm_value,
        )

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

    def train(self) -> None:
        """Run the configured global optimizer steps by Dataset epoch."""
        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.base.state.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__ = ["TextTrainer"]