import argparse
import json

from autotrain.trainers.common import monitor
from autotrain.trainers.vlm import utils
from autotrain.trainers.vlm.params import VLMTrainingParams


def parse_args():
    # get training_config.json from the end user
    parser = argparse.ArgumentParser()
    parser.add_argument("--training_config", type=str, required=True)
    return parser.parse_args()


@monitor
def train(config):
    if isinstance(config, dict):
        config = VLMTrainingParams(**config)

    if not utils.check_model_support(config):
        raise ValueError(f"model `{config.model}` not supported")

    if config.trainer in ("vqa", "captioning"):
        from autotrain.trainers.vlm.train_vlm_generic import train as train_generic

        train_generic(config)

    else:
        raise ValueError(f"trainer `{config.trainer}` not supported")


if __name__ == "__main__":
    # Check if MPS should be disabled before importing torch-dependent modules
    import os

    if os.environ.get("AUTOTRAIN_DISABLE_MPS") == "1":
        # Monkey-patch torch to disable MPS detection
        import torch

        original_is_available = torch.backends.mps.is_available
        torch.backends.mps.is_available = lambda: False
        # Also prevent MPS from being used accidentally
        if hasattr(torch.backends.mps, "is_built"):
            torch.backends.mps.is_built = lambda: False

    _args = parse_args()
    training_config = json.load(open(_args.training_config))
    _config = VLMTrainingParams(**training_config)
    train(_config)
