import os
import sys
from typing import Tuple
from utils.base_arguments import BaseArguments
from utils.argument_parser import CustomHFArgumentParser as ArgParser
from utils.pruning_utils import PruningArguments
from utils.dataset_utils import CalibrationDatasetArguments
from utils.model_utils import ModelArguments
from utils.evaluation_arguments import EvaluationArguments


def parse_config() -> Tuple[BaseArguments, ModelArguments, CalibrationDatasetArguments, PruningArguments, EvaluationArguments]:
    parser = ArgParser([BaseArguments, ModelArguments, CalibrationDatasetArguments, PruningArguments, EvaluationArguments])
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        base_args, model_args, data_args, pruning_args, evaluation_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
        base_args, model_args, data_args, pruning_args, evaluation_args = parser.parse_yaml_file(yaml_file=os.path.abspath(sys.argv[1]))
    else:
        base_args, model_args, data_args, pruning_args, evaluation_args = parser.parse_args_into_dataclasses()
    return (
        BaseArguments(**vars(base_args)),
        ModelArguments(**vars(model_args)),
        CalibrationDatasetArguments(**vars(data_args)),
        PruningArguments(**vars(pruning_args)),
        EvaluationArguments(**vars(evaluation_args))
    )