from __future__ import annotations

from datahub.db import HubDB


def resolve_code(db: HubDB, raw: str) -> str | None:
    text = str(raw or "").strip().upper()
    if not text:
        return None
    if "." in text:
        row = db.fetchone("SELECT ts_code FROM stock_master WHERE ts_code = ?", (text,))
        if row:
            return row["ts_code"]
        # indices are not always in stock_master
        return text
    matches = db.fetchall(
        "SELECT ts_code FROM stock_master WHERE symbol = ? OR ts_code LIKE ?",
        (text, f"{text}.%"),
    )
    if len(matches) == 1:
        return matches[0]["ts_code"]
    if len(matches) > 1:
        return None
    # unique exchange guess for 6-digit codes
    suffix = "SH" if text.startswith("6") or text.startswith("9") else "SZ" if text.startswith(("0", "3")) else "BJ"
    return f"{text}.{suffix}"