path: root/worker/data.py

"""AKShare data adapter plus bounded provider suggestion search.

No silent provider/adjustment fallbacks. No synthetic/fabricated data.
"""
from __future__ import annotations

import json
import re
from typing import Any

import akshare as ak  # verified import at module load
import pandas as pd
import requests

from .normalize import COLUMN_MAP, map_columns

SUPPORTED_ASSET_TYPES = {"stock", "etf", "index"}
SUPPORTED_FREQUENCIES = {"daily"}
SUPPORTED_ADJUSTMENTS = {"none", "qfq", "hfq"}
IDENTITY_EXCHANGES = {"SH", "SZ", "BJ"}


SUPPORTED_SOURCES = {"eastmoney", "tencent", "sina", "auto"}
_REQUEST_TIMEOUT_SECS = 15.0


class DataError(ValueError):
    def __init__(self, code: str, message: str, details: Any = None):
        super().__init__(message)
        self.code = code
        self.details = details


class FetchResult(tuple):
    """(normalized_df, endpoint, params) with .provider and .raw extras.

    Cache identity and manifest MUST use the actual provider/endpoint that
    served the data (auto fallback records it here).
    """

    def __new__(cls, df, endpoint, params, raw, provider):
        obj = super().__new__(cls, (df, endpoint, params))
        obj.raw = raw
        obj.provider = provider
        return obj

    @property
    def df(self):
        return self[0]

    @property
    def endpoint(self):
        return self[1]

    @property
    def params(self):
        return self[2]


def split_identity(instrument: dict) -> tuple[str, str]:
    """-> (exchange, code). Accepts 'SH#600000' or bare code with market identity."""
    symbol = str(instrument["symbol"])
    if "#" in symbol:
        exch, code = symbol.split("#", 1)
        if exch not in IDENTITY_EXCHANGES:
            raise DataError("bad_identity", f"unknown exchange prefix: {exch}")
    else:
        code = symbol
        market = instrument.get("market", "cn")
        if market != "cn":
            raise DataError("unsupported_market", f"market not supported by this adapter: {market}")
        if len(code) == 6 and code[0] in "369":
            exch = "SH"
        else:
            exch = "SZ"
    if not code.isdigit() or len(code) != 6:
        raise DataError("bad_identity", "A-share code must be a 6-digit number")
    return exch, code


def _fetch_eastmoney(instrument: dict, start: str, end: str,
                     frequency: str, adjustment: str) -> FetchResult:
    exch, code = split_identity(instrument)
    s, e = start.replace("-", ""), end.replace("-", "")
    asset = instrument["asset_type"]
    if asset == "stock":
        endpoint = "stock_zh_a_hist"
        params = {"symbol": code, "period": "daily", "start_date": s, "end_date": e,
                  "adjust": "" if adjustment == "none" else adjustment}
        raw = ak.stock_zh_a_hist(**params)
    elif asset == "etf":
        endpoint = "fund_etf_hist_em"
        params = {"symbol": code, "period": "daily", "start_date": s, "end_date": e,
                  "adjust": "" if adjustment == "none" else adjustment}
        raw = ak.fund_etf_hist_em(**params)
    else:
        endpoint = "index_zh_a_hist"
        params = {"symbol": code, "period": "daily", "start_date": s, "end_date": e}
        raw = ak.index_zh_a_hist(**params)
    if not isinstance(raw, pd.DataFrame):
        raise DataError("bad_provider_response", f"{endpoint} did not return a DataFrame")
    df = _to_canonical(raw, f"{exch}#{code}")
    return FetchResult(df, endpoint, params, raw, "eastmoney")


def _fetch_tencent(instrument: dict, start: str, end: str,
                   frequency: str, adjustment: str) -> FetchResult:
    exch, code = split_identity(instrument)
    asset = instrument["asset_type"]
    s, e = start.replace("-", ""), end.replace("-", "")
    tx_symbol = f"{exch.lower()}{code}"
    if asset == "stock":
        endpoint = "stock_zh_a_hist_tx"
        params = {"symbol": tx_symbol, "start_date": s, "end_date": e,
                  "adjust": "" if adjustment == "none" else adjustment}
        raw = ak.stock_zh_a_hist_tx(**params, timeout=_REQUEST_TIMEOUT_SECS)
    elif asset == "index":
        if frequency != "daily":
            raise DataError("unsupported_frequency", "tencent source supports daily only")
        endpoint = "stock_zh_index_daily_tx"
        params = {"symbol": tx_symbol, "start_date": s, "end_date": e}
        raw = ak.stock_zh_index_daily_tx(**params)
    else:
        raise DataError("source_unavailable", "tencent source has no listed-ETF daily adapter")
    if not isinstance(raw, pd.DataFrame):
        raise DataError("bad_provider_response", f"{endpoint} did not return a DataFrame")
    symbol = f"{exch}#{code}"
    df = _to_canonical(raw, symbol)
    warnings = []
    if asset == "index":
        warnings.append("tencent index data is labeled 前复权 by the provider; "
                        "adjustment-factors mixed across sources are not supported")
    df.attrs["source_warnings"] = warnings
    return FetchResult(df, endpoint, params, raw, "tencent")


def _fetch_sina(instrument: dict, start: str, end: str,
                frequency: str, adjustment: str) -> FetchResult:
    exch, code = split_identity(instrument)
    asset = instrument["asset_type"]
    if frequency != "daily":
        raise DataError("unsupported_frequency", "sina source supports daily only")
    if asset != "etf":
        raise DataError("source_unavailable",
                        f"sina source has no {asset} daily adapter in this worker")
    if adjustment != "none":
        raise DataError("unsupported_adjustment",
                        "sina ETF klines are unadjusted only (no adjust parameter); "
                        f"adjustment '{adjustment}' cannot be served by sina")
    endpoint = "fund_etf_hist_sina"
    params = {"symbol": f"{exch.lower()}{code}"}
    raw = ak.fund_etf_hist_sina(**params)
    if not isinstance(raw, pd.DataFrame):
        raise DataError("bad_provider_response", f"{endpoint} did not return a DataFrame")
    df = _to_canonical(raw, f"{exch}#{code}")
    # sina has no date-range parameter: slice the full history locally to the
    # requested window (ISO strings compare lexicographically like dates)
    df = df[(df["date"] >= start) & (df["date"] <= end)].reset_index(drop=True)
    df.attrs["source_warnings"] = [
        "sina returns the full trading history without date parameters; sliced "
        "locally to the requested range; data is unadjusted (no adjust parameter)",
        "sina volume unit is 股 (shares); verified to be 100x the 手 (lots) "
        "convention used by lot-based providers on the same session "
        "(2026-09-17 evidence, docs/recovery-01-plan.md)",
    ]
    return FetchResult(df, endpoint, params, raw, "sina")


def _to_canonical(raw: pd.DataFrame, canonical_symbol: str) -> pd.DataFrame:
    df = map_columns(raw.copy())
    if not df.empty and "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d")
    df["symbol"] = canonical_symbol
    return df


def fetch_source(instrument: dict, start: str, end: str, frequency: str,
                 adjustment: str, fields: list[str] | None = None,
                 source: str = "auto") -> FetchResult:
    if frequency not in SUPPORTED_FREQUENCIES:
        raise DataError("unsupported_frequency", f"frequency '{frequency}' not supported; supported: daily")
    if adjustment not in SUPPORTED_ADJUSTMENTS:
        raise DataError("unsupported_adjustment", f"adjustment '{adjustment}' not supported")
    if instrument.get("asset_type") not in {"stock", "etf", "index"}:
        raise DataError("unsupported_asset_type", f"asset_type '{instrument.get('asset_type')}' not supported")
    if source not in SUPPORTED_SOURCES:
        raise DataError("bad_source", f"source '{source}' not supported; supported: {sorted(SUPPORTED_SOURCES)}")
    if instrument["asset_type"] == "index" and adjustment != "none":
        raise DataError("unsupported_adjustment", "indexes have no adjustment factors; adjustment must be 'none'")

    def run(provider: str) -> FetchResult:
        if provider == "eastmoney":
            return _fetch_eastmoney(instrument, start, end, frequency, adjustment)
        if provider == "sina":
            return _fetch_sina(instrument, start, end, frequency, adjustment)
        return _fetch_tencent(instrument, start, end, frequency, adjustment)

    if source != "auto":
        return run(source)
    # auto: eastmoney first, transparent same-symbol fallback, honestly labeled.
    # tencent has no listed-ETF daily adapter, so ETFs fall back to sina.
    fallback = "sina" if instrument["asset_type"] == "etf" else "tencent"
    try:
        return run("eastmoney")
    except Exception as em_err:
        try:
            res = run(fallback)
        except Exception as fb_err:
            raise DataError(
                "provider_unavailable",
                f"eastmoney failed ({em_err}); {fallback} fallback failed ({fb_err})")
        res.df.attrs["source_warnings"] = list(getattr(res.df, "attrs", {}).get("source_warnings", [])) + [
            f"provider_fallback: eastmoney attempt failed ({type(em_err).__name__}); "
            f"served by {fallback} for the SAME symbol; sources may differ in "
            "adjustment method and units (check provider warnings)"
        ]
        return res


def fetch_source_with_warnings(instrument: dict, start: str, end: str, frequency: str,
                               adjustment: str, fields: list[str] | None = None,
                               source: str = "auto"):
    """fetch_source plus source warnings pulled off the returned frame."""
    res = fetch_source(instrument, start, end, frequency, adjustment, fields, source)
    src_warnings = list(res.df.attrs.get("source_warnings", []))
    return res, src_warnings


# ---- provider suggestion search (bounded, direct HTTP, no full catalogs) ----
#
# Root-cause fix for the production search stall: the legacy path fetched the
# ENTIRE stock catalog and the ENTIRE ETF snapshot on every query. Instead we
# hit the actual suggestion endpoints of the same data providers used by
# fetch — bounded HTTP calls with short timeouts, then validated identity
# classification per item. Asset class comes ONLY from the provider's own
# classification field (Eastmoney `Classify` / Tencent hint tail token),
# never guessed from the numeric code shape.

SUGGEST_TIMEOUT_SECS = 4.0  # per-source hard timeout, well under overall budget

_EASTMONEY_URL = "https://searchapi.eastmoney.com/api/suggest/get"
_TENCENT_URL = "https://smartbox.gtimg.cn/s3/"

# Live-probed provider-stated classes (2026-09, see docs/search-fix.md):
# Eastmoney Classify: AStock→stock, Index→index, Fund→fund (ambiguous:
# same SecurityType/Classify covers ETF AND LOF), Bond/HK/OTCFUND/... unsupported.
# Tencent hint tail: GP-A→stock, ETF→etf, ZS→index, LOF/others unsupported.
_EASTMONEY_CLASSIFY = {"AStock": "stock", "Index": "index", "Fund": "etf"}
_TENCENT_TAGS = {"GP-A": "stock", "ETF": "etf", "ZS": "index"}
_MKTNUM_EXCHANGE = {"1": "SH", "0": "SZ"}
_TENCENT_EXCHANGE = {"sh": "SH", "sz": "SZ"}

_HEADERS = {"User-Agent": "Mozilla/5.0 (strategy-lab instrument search)"}


def _http_json(url: str, params: dict) -> dict:
    r = requests.get(url, params=params, timeout=SUGGEST_TIMEOUT_SECS, headers=_HEADERS)
    r.raise_for_status()
    return r.json()


def _http_text(url: str, params: dict) -> str:
    r = requests.get(url, params=params, timeout=SUGGEST_TIMEOUT_SECS, headers=_HEADERS)
    r.raise_for_status()
    return r.text


def _search_eastmoney(q: str, limit: int) -> list[dict]:
    """Bounded direct eastmoney suggest call. Returns validated items only."""
    payload = _http_json(
        _EASTMONEY_URL,
        {"input": q, "type": 14, "count": max(5, min(limit, 100))},
    )
    table = payload.get("QuotationCodeTable") or {}
    if table.get("Status") != 0:
        raise RuntimeError(f"eastmoney suggest status={table.get('Status')!r}")
    rows = table.get("Data") or []
    items = []
    for row in rows[:limit]:
        code = str(row.get("Code", "")).strip()
        name = str(row.get("Name", "")).strip()
        classify = str(row.get("Classify", "")).strip()
        asset = _EASTMONEY_CLASSIFY.get(classify)
        exchange = _MKTNUM_EXCHANGE.get(str(row.get("MktNum", "")).strip())
        # identity must be fully validated by the provider reply itself; the
        # asset class is provider-stated, never inferred from the code digits
        if not code or not exchange or asset is None:
            continue
        items.append({
            "symbol": code,
            "canonical_symbol": f"{exchange}#{code}",
            "market": "cn",
            "asset_type": asset,
            "name": name,
            "currency": "CNY",
            "source": "eastmoney",
        })
    return items


def _search_tencent(q: str, limit: int) -> list[dict]:
    """Bounded direct tencent smartbox call. Returns validated items only."""
    text = _http_text(_TENCENT_URL, {"q": q, "t": "all"})
    m = re.search(r'v_hint="(.*)"', text)
    if not m:
        return []  # empty suggestion is a valid provider-no-match reply
    body = m.group(1)
    # provider encodes non-ASCII as \uXXXX escapes; decode JSON-style safely
    try:
        body = json.loads(f'"{body}"')
    except ValueError:
        body = decode_unicode_escapes(body)
    items = []
    for entry in body.split("^"):
        parts = entry.strip().split("~")
        if len(parts) < 5:
            continue
        mkt, code, name, _py, tag = parts[0].lower(), parts[1], parts[2], parts[3], parts[4].strip()
        asset = _TENCENT_TAGS.get(tag)
        exchange = _TENCENT_EXCHANGE.get(mkt)
        # class and exchange both provider-stated; LOF/bonds/bj are excluded
        if not code or not exchange or asset is None:
            continue
        items.append({
            "symbol": code,
            "canonical_symbol": f"{exchange}#{code}",
            "market": "cn",
            "asset_type": asset,
            "name": name,
            "currency": "CNY",
            "source": "tencent",
        })
    return items[:limit]


def decode_unicode_escapes(s: str) -> str:
    return re.sub(r"\\u([0-9a-fA-F]{4})", lambda m: chr(int(m.group(1), 16)), s)


def search_instruments(q: str, limit: int = 50) -> dict:
    """Bounded direct provider suggestion search.

    Per-source HTTP timeout 4s; one source failing is visible but not fatal if
    the other serves results. Both failing -> honest failed status, never an
    empty success. Each item carries its actual provenance source.
    """
    out = {"source": "provider_suggest", "status": "failed",
           "items": [], "error": None, "providers": {}}
    errors: list[str] = []
    sources_used: list[str] = []
    results: dict[str, list[dict]] = {}

    try:
        results["eastmoney"] = _search_eastmoney(q, limit)
        sources_used.append("eastmoney")
        out["providers"]["eastmoney"] = "ok"
    except Exception as exc:
        errors.append(f"eastmoney: {type(exc).__name__}: {exc}")
        out["providers"]["eastmoney"] = f"failed: {type(exc).__name__}"

    try:
        results["tencent"] = _search_tencent(q, limit)
        sources_used.append("tencent")
        out["providers"]["tencent"] = "ok"
    except Exception as exc:
        errors.append(f"tencent: {type(exc).__name__}: {exc}")
        out["providers"]["tencent"] = f"failed: {type(exc).__name__}"

    if not sources_used:
        out["error"] = {"code": "providers_unavailable",
                        "message": "all providers failed: " + "; ".join(errors)}
        return out

    tencent_by_code = {it["symbol"]: it for it in results.get("tencent", [])}
    tencent_ok = "tencent" in results
    seen: set[tuple[str, str, str]] = set()
    items = []
    for provider in ("eastmoney", "tencent"):
        for it in results.get(provider, []):
            # Eastmoney's Classify "Fund" is ambiguous (verified live: ETF and
            # LOF share the same Classify/SecurityType). Only emit Eastmoney
            # fund items when Tencent — whose tags distinguish ETF from LOF —
            # serves the same code; otherwise the LOF guarantee cannot hold.
            if provider == "eastmoney" and it["asset_type"] == "etf":
                cross = tencent_by_code.get(it["symbol"])
                if not (tencent_ok and cross and cross["asset_type"] == "etf"):
                    continue
            key = (it["canonical_symbol"], it["asset_type"], it["source"])
            if key in seen:
                continue
            seen.add(key)
            items.append(it)
            if len(items) >= limit:
                break
        if len(items) >= limit:
            break

    out["status"] = "ready"
    out["items"] = items
    out["source"] = "provider_suggest"
    out["error"] = ({"code": "partial_providers",
                     "message": "some providers failed: " + "; ".join(errors)}) if errors else None
    if errors:
        out["warnings"] = errors
    return out