"""A-share stock-code normalization helpers."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class NormalizedStockCode:
    code: str
    market: str
    suffix: str
    prefixed: str
    display: str


class InvalidSymbolError(ValueError):
    def __init__(self, symbol: str, message: str | None = None) -> None:
        self.symbol = symbol
        self.code = "invalid_symbol"
        super().__init__(
            message
            or f"invalid_symbol: {symbol} is not a valid A-share stock code; "
            "ETF/fund codes should use ETF or fund endpoints"
        )


def normalize_stock_code(stock_code: str) -> NormalizedStockCode:
    raw = str(stock_code or "").strip()
    if not raw:
        raise InvalidSymbolError(stock_code, "invalid_symbol: stock_code is required")

    lower = raw.lower()
    if lower.startswith(("sh", "sz", "bj")) and len(lower) >= 8:
        market = lower[:2]
        code = lower[2:8]
    elif "." in raw:
        code_part, suffix_part = raw.split(".", 1)
        code = code_part.strip()
        market = suffix_part.strip().lower()
    else:
        code = raw[:6]
        inferred = _infer_a_share_market(code)
        if inferred:
            market = inferred
        elif code.startswith(("8", "4", "920")):
            market = "bj"
        elif code.startswith(("510", "512", "513", "515", "588", "159")):
            market = "sz" if code.startswith("159") else "sh"
        elif code.startswith(("600", "601", "603", "605", "688")):
            market = "sh"
        else:
            market = "sz"

    if len(code) != 6 or not code.isdigit():
        raise InvalidSymbolError(stock_code)
    if market not in {"sh", "sz", "bj"}:
        raise InvalidSymbolError(stock_code, f"invalid_symbol: unsupported market suffix for {stock_code}")
    expected_market = _infer_a_share_market(code)
    if expected_market is None:
        raise InvalidSymbolError(stock_code)
    if market != expected_market:
        raise InvalidSymbolError(
            stock_code,
            f"invalid_symbol: {stock_code} does not match expected {expected_market.upper()} market for A-share stocks",
        )

    suffix = market.upper()
    return NormalizedStockCode(
        code=code,
        market=market,
        suffix=suffix,
        prefixed=f"{market}{code}",
        display=f"{code}.{suffix}",
    )


def _infer_a_share_market(code: str) -> str | None:
    if len(code) != 6 or not code.isdigit():
        return None
    if code.startswith(("600", "601", "603", "605", "688")):
        return "sh"
    if code.startswith(("000", "001", "002", "003", "300", "301")):
        return "sz"
    if code.startswith(("920", "8", "4")):
        return "bj"
    return None