import os
from typing import Union


def _maybestr2bool(value: Union[bool, str], error: str) -> bool:
    if isinstance(value, bool):
        return value
    elif isinstance(value, str):
        if value.lower() in ("0", "false"):
            return False
        elif value.lower() in ("1", "true"):
            return True
        else:
            raise ValueError(error)
    else:
        raise ValueError(error)


class _JaxtypingConfig:
    def __init__(self):
        self.update("jaxtyping_disable", os.environ.get("JAXTYPING_DISABLE", "0"))
        self.update(
            "jaxtyping_remove_typechecker_stack",
            os.environ.get("JAXTYPING_REMOVE_TYPECHECKER_STACK", "0"),
        )

    def update(self, item: str, value):
        if item.lower() == "jaxtyping_disable":
            msg = (
                "Unrecognised value for `JAXTYPING_DISABLE`. Valid values are "
                "`JAXTYPING_DISABLE=0` (the default) or `JAXTYPING_DISABLE=1` (to "
                "disable runtime type checking)."
            )
            self.jaxtyping_disable = _maybestr2bool(value, msg)
        elif item.lower() == "jaxtyping_remove_typechecker_stack":
            msg = (
                "Unrecognised value for `JAXTYPING_REMOVE_TYPECHECKER_STACK`. Valid "
                "values are `JAXTYPING_REMOVE_TYPECHECKER_STACK=0` (the default) or "
                "`JAXTYPING_REMOVE_TYPECHECKER_STACK=1` (to remove the stack frames "
                "from the typechecker in `jaxtyped(typechecker=...)`, when it raises a "
                "runtime type-checking error)."
            )
            self.jaxtyping_remove_typechecker_stack = _maybestr2bool(value, msg)
        else:
            raise ValueError(f"Unrecognised config value {item}")


config = _JaxtypingConfig()