from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass

from .constants import DISABLE_COLOR_ENV_VARS
from .utils import get_env_value

# Toggled internally to suppress colors for a block of output. A ContextVar is
# used instead of mutating os.environ, which is not thread safe and can segfault
# unrelated threads that read the environment from native code.
_disable_color: ContextVar[bool] = ContextVar("syrupy_disable_color", default=False)


@contextmanager
def disable_color() -> Iterator[None]:
    """Suppress terminal colors for the duration of the context."""
    token = _disable_color.set(True)
    try:
        yield
    finally:
        _disable_color.reset(token)


def _is_color_disabled() -> bool:
    return _disable_color.get() or any(map(get_env_value, DISABLE_COLOR_ENV_VARS))


@dataclass
class TerminalCodes:
    ESC: str = "\x1b["
    END: str = "m"
    FOREGROUND_256: str = f"{ESC}38;5;"
    BACKGROUND_256: str = f"{ESC}48;5;"

    STYLES = {
        "bold": "1",
        "dim": "2",
        "italic": "3",
        "underline": "4",
        "reset": "0",
    }

    COLORS = {
        "black": "0",
        "red": "1",
        "green": "2",
        "yellow": "3",
    }


def _attr(style: str) -> str:
    if _is_color_disabled():
        return ""
    return f"{TerminalCodes.ESC}{TerminalCodes.STYLES[style]}{TerminalCodes.END}"


def _fg(color: int | str) -> str:
    if _is_color_disabled():
        return ""
    color_code = TerminalCodes.COLORS[color] if isinstance(color, (str,)) else color
    return f"{TerminalCodes.FOREGROUND_256}{str(color_code)}{TerminalCodes.END}"


def _bg(color: int) -> str:
    if _is_color_disabled():
        return ""
    return f"{TerminalCodes.BACKGROUND_256}{str(color)}{TerminalCodes.END}"


def _stylize(text: str | int, formatting: str) -> str:
    if _is_color_disabled():
        return str(text)
    return f"{formatting}{text}{_attr('reset')}"


def reset(text: str | int) -> str:
    return _stylize(text, _attr("reset"))


def red(text: str | int) -> str:
    return _stylize(text, _fg("red"))


def yellow(text: str | int) -> str:
    return _stylize(text, _fg("yellow"))


def green(text: str | int) -> str:
    return _stylize(text, _fg("green"))


def bold(text: str | int) -> str:
    return _stylize(text, _attr("bold"))


def error_style(text: str | int) -> str:
    return bold(red(text))


def warning_style(text: str | int) -> str:
    return bold(yellow(text))


def success_style(text: str | int) -> str:
    return bold(green(text))


def snapshot_style(text: str | int) -> str:
    return _stylize(text, _bg(225) + _fg(90))


def snapshot_diff_style(text: str | int) -> str:
    return _stylize(text, _bg(90) + _fg(225))


def received_style(text: str | int) -> str:
    return _stylize(text, _bg(195) + _fg(23))


def received_diff_style(text: str | int) -> str:
    return _stylize(text, _bg(23) + _fg(195))


def context_style(text: str | int) -> str:
    return _stylize(text, _attr("dim"))
