import os
import shutil
import time

from transformers import TrainerCallback, TrainerState, TrainerControl
from emissary_ml_toolkit.utils.logging import get_logger
from emissary_ml_toolkit.webhooks.client import WebhookClient
from emissary_ml_toolkit.webhooks.protocol import TrainingProgress
from emissary_ml_toolkit.storage.base import BaseStorage
from emissary_ml_toolkit.config.schemas import TrainConfig


logger = get_logger(__name__)


class EmissaryTrainerCallback(TrainerCallback):
    def __init__(self,
                 storage: BaseStorage,
                 webhook_client: WebhookClient,
                 train_config: TrainConfig,):
        self.webhook_client = webhook_client
        self.storage = storage
        self.train_config = train_config

        self.current_epoch = 0
        self.start_time = None

    def on_save(self, args, state: TrainerState, control: TrainerControl, **kwargs):
        if not state.is_world_process_zero:
            # is_world_process_zero
            # Whether or not this process is the global main process
            # (When training in a distributed fashion on several machines,
            # this is only going to be `True` for one process).
            return

        if self.current_epoch >= state.num_train_epochs:
            return

        self.current_epoch += 1
        logger.info(f"Saving checkpoint {self.current_epoch} to the storage")
        checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")

        embeddings_source = os.path.join("checkpoints", "new_embeddings.pt")
        if os.path.exists(embeddings_source):
            embeddings_dest = os.path.join(checkpoint_dir, "new_embeddings.pt")
            shutil.copy2(embeddings_source, embeddings_dest)

        task_details_source = "task_details.json"
        if os.path.exists(task_details_source):
            task_details_dest = os.path.join(checkpoint_dir, "task_details.json")
            shutil.copy2(task_details_source, task_details_dest)

        classification_head_weights_source = "classification_tensors.pt"
        if os.path.exists(classification_head_weights_source):
            classification_head_weights_dest = os.path.join(checkpoint_dir, "classification_tensors.pt")
            shutil.copy2(classification_head_weights_source, classification_head_weights_dest)

        zip_artifacts = self.storage.zip_directory(checkpoint_dir, "artifacts.zip")
        storage_path = f"{self.train_config.training_output.base_path}/{str(self.current_epoch)}/artifacts.zip"
        self.storage.upload_file(zip_artifacts, storage_path)

        logger.info(f"Checkpoint {self.current_epoch} saved")
        self.webhook_client.send_checkpoint_upload(self.current_epoch)

    def on_train_begin(self, args, state: TrainerState, control: TrainerControl, **kwargs):
        logger.info("Starting training loop")
        self.start_time = time.time()
        self.webhook_client.send_training_status(status="Running")

    def on_log(self, args, state: TrainerState, control: TrainerControl, logs=None, **kwargs):
        current_step = state.global_step
        total_steps = state.max_steps

        current_epoch = int(state.epoch + 1 or 0)
        total_epochs = int(getattr(args, "num_train_epochs", 0) or 0)

        # Calculate ETA
        eta_seconds = None
        if current_step > 0:
            elapsed_time = time.time() - self.start_time
            avg_time_per_step = elapsed_time / current_step
            steps_remaining = total_steps - current_step
            eta_seconds = int(avg_time_per_step * steps_remaining)

        training_progress = TrainingProgress(
            current_epoch=current_epoch,
            current_step=current_step,
            total_epoch=total_epochs,
            total_step=total_steps,
            estimated_remaining_time=eta_seconds,
        )

        logger.info(f"Total Epoch/Step: {total_epochs} / {total_steps}; Current Epoch/Step: {current_epoch} / {current_step}")

        train_loss = None

        if logs is not None and "loss" in logs:
            train_loss = float(logs["loss"])

        logger.info(f"Loss at step: {current_step} in epoch: {current_epoch}: {train_loss}")
        self.webhook_client.send_training_progress(training_progress, train_loss)
