from ._types import EquityTickerClassStyle, TickerNormalized, TickerUpperCase

# Common index tickers that should be prefixed with $
DEFAULT_DOLLAR_PREFIX_TICKERS: set[str] = {
    "SPX",
    "VIX",
    "NDX",
    "RUT",
    "DJI",
    "DJX",
    "VVIX",
    "VIX9D",
    "VIX3M",
    "VIX6M",
    "VIX1Y",
    "VIXEQ",
    "VIX1D",
    "VXN",
    "RVX",
    "VXD",
}

# Module-level set that users can extend
dollar_prefix_tickers: set[str] = DEFAULT_DOLLAR_PREFIX_TICKERS.copy()


def add_dollar_prefix_tickers(*tickers: str) -> None:
    """Add tickers to the set that should be prefixed with $."""
    for t in tickers:
        dollar_prefix_tickers.add(t.upper().lstrip("$"))


def normalize_ticker(
    ticker: str,
    prefix_dollar: bool = True,
    normalize_index_tickers: bool = True,
    equity_ticker_class_style: EquityTickerClassStyle = "/",
    prefix_futures_ticker: bool = False,
) -> TickerNormalized | TickerUpperCase:
    """
    Normalize a ticker symbol.

    - Uppercase the ticker
    - Optionally add $ prefix for index tickers
    - Normalize index ticker variants (SPXW -> SPX, etc.)
    - Handle equity ticker class styles (BRK.B vs BRK/B)

    Args:
        ticker: The ticker symbol to normalize
        prefix_dollar: If True, add $ prefix for index tickers
        normalize_index_tickers: If True, normalize variants like SPXW -> SPX
        equity_ticker_class_style: "/" or "." for class shares (e.g., BRK/B or BRK.B)
        prefix_futures_ticker: If True, prefix with / for futures

    Returns:
        Normalized ticker string
    """
    ticker = ticker.upper().lstrip("/$")

    if equity_ticker_class_style == "/":
        ticker = ticker.replace(".", "/")
    elif equity_ticker_class_style == ".":
        ticker = ticker.replace("/", ".")

    if normalize_index_tickers:
        if ticker == "SPXW":
            ticker = "SPX"
        elif ticker == "NDXP":
            ticker = "NDX"
        elif ticker == "RUTW":
            ticker = "RUT"
        elif ticker == "VIXW":
            ticker = "VIX"

    if prefix_dollar:
        if ticker in dollar_prefix_tickers:
            return f"${ticker}"
        else:
            return ticker
    elif prefix_futures_ticker:
        return f"/{ticker.strip('$/')}"
    else:
        return ticker.removeprefix("$")
