# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates

#  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.

import argparse
import json
import os
from pathlib import Path
from typing import Union

import toml
import yaml


class ArgumentParser(argparse.ArgumentParser):
    def __init__(self, use_config: bool = True, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if use_config:
            self.add_argument(
                "--config",
                type=Path,
                help="Path to the config file",
                default=None,
            )

    def parse_args(
        self, use_config: bool = True, *args, **kwargs
    ) -> argparse.Namespace:
        args = super().parse_args(*args, **kwargs)

        if use_config:
            if args.config:
                self.load_from_config_file(args, args.config)

            self.load_from_env_variables(args)

        return args

    def load_from_config_file(
        self,
        args: argparse.Namespace,
        file_path: Path,
        strict: bool = True,
    ):
        if file_path.suffix == ".json":
            with open(file_path, "r", encoding="utf-8") as f:
                config = json.load(f)
        elif file_path.suffix == ".toml":
            with open(file_path, "r", encoding="utf-8") as f:
                config = toml.load(f)
        elif file_path.suffix in [".yaml", ".yml"]:
            with open(file_path, "r", encoding="utf-8") as f:
                config = yaml.safe_load(f)
        else:
            raise ValueError(
                "Unsupported config file format. Use .json, .toml, or .yaml/.yml"
            )

        for key, value in config.items():
            if strict and not hasattr(args, key):
                raise ValueError(
                    f"Found invalid key `{key}` in config file: {file_path}"
                )

            if getattr(args, key) == self.get_default(key):
                setattr(args, key, value)

    def load_from_env_variables(self, args: argparse.Namespace):
        for key in vars(args).keys():
            env_key = key.replace("-", "_").upper()
            env_value = os.getenv(env_key)
            if env_value is not None and getattr(args, key) == self.get_default(key):
                setattr(args, key, env_value)

    def bool(self, value: Union[str, bool, None]) -> bool:
        if isinstance(value, bool):
            return value

        if not isinstance(value, str):
            raise argparse.ArgumentTypeError("Boolean value expected.")

        if value.lower() in {"true", "yes", "1"}:
            return True
        elif value.lower() in {"false", "no", "0"}:
            return False
        else:
            raise argparse.ArgumentTypeError(f"Boolean expected, got `{value}`")
